support cfg_zero*

app.py

@@ -28,7 +28,7 @@ def clear_cache(transformer):
 
 # Define the Gradio interface
 @spaces.GPU()
-def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type):
+def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, use_zero_init, zero_steps):
     # Set the control type
     if control_type == "Ghibli":
         lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
@@ -47,6 +47,8 @@ def single_condition_generate_image(prompt, spatial_img, height, width, seed, co
         subject_images=[],
         spatial_images=spatial_imgs,
         cond_size=512,
+        use_zero_init=use_zero_init,
+        zero_steps=int(zero_steps)
     ).images[0]
     clear_cache(pipe.transformer)
     return image
@@ -86,13 +88,16 @@ with gr.Blocks() as demo:
                 seed = gr.Number(label="Seed", value=42)
                 control_type = gr.Dropdown(choices=control_types, label="Control Type")
                 single_generate_btn = gr.Button("Generate Image")
+                use_zero_init = gr.Checkbox(label="Use zero init", default=False)
+                zero_steps = gr.Number(label="Zero Init Steps", value=1)
+                
             with gr.Column():
                 single_output_image = gr.Image(label="Generated Image")
 
         # Add examples for Single Condition Generation
         gr.Examples(
             examples=single_examples,
-            inputs=[prompt, spatial_img, height, width, seed, control_type],
+            inputs=[prompt, spatial_img, height, width, seed, control_type, use_zero_init, zero_steps],
             outputs=single_output_image,
             fn=single_condition_generate_image,
             cache_examples=False,  # 缓存示例结果以加快加载速度