init

app.py

@@ -9,7 +9,10 @@ from pipeline_freescale import StableDiffusionXLPipeline
 from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 
 @spaces.GPU(duration=120)
-def infer_gpu_part(pipe, generator, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale):
+def infer_gpu_part(pipe, generator, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu):
+    if not disable_freeu:
+        register_free_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
+        register_free_crossattn_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
     pipe = pipe.to("cuda")
     generator = generator.to("cuda")
     resul = pipe(prompt, negative_prompt=negative_prompt, generator=generator,
@@ -40,14 +43,11 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
 
     model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
     pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=torch.float16)
-    if not disable_freeu:
-        register_free_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
-        register_free_crossattn_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
-
+    
     generator = torch.Generator()
     generator = generator.manual_seed(seed)
 
-    result = infer_gpu_part(pipe, generator, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale)
+    result = infer_gpu_part(pipe, generator, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu)
 
     image = result.images[0]
     save_path = 'output.png'
@@ -204,4 +204,5 @@ with gr.Blocks(css=css) as demo:
             outputs=[image_result],
             api_name="freescalehf")
 
-demo.queue(max_size=8).launch(show_api=True)
+if __name__ == "__main__":
+    demo.queue(max_size=8).launch(show_api=True)