init

app.py

@@ -9,10 +9,41 @@ 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, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale):
+def infer_gpu_fast(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu):
     pipe = pipe.to("cuda") 
     generator = torch.Generator(device='cuda')
     generator = generator.manual_seed(seed)
+    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)
+    result = pipe(prompt, negative_prompt=negative_prompt, generator=generator,
+                num_inference_steps=ddim_steps, guidance_scale=guidance_scale,
+                resolutions_list=resolutions_list, fast_mode=fast_mode, cosine_scale=cosine_scale,
+                )
+    return result
+
+@spaces.GPU(duration=240)
+def infer_gpu_mid(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu):
+    pipe = pipe.to("cuda") 
+    generator = torch.Generator(device='cuda')
+    generator = generator.manual_seed(seed)
+    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)
+    result = pipe(prompt, negative_prompt=negative_prompt, generator=generator,
+                num_inference_steps=ddim_steps, guidance_scale=guidance_scale,
+                resolutions_list=resolutions_list, fast_mode=fast_mode, cosine_scale=cosine_scale,
+                )
+    return result
+
+@spaces.GPU(duration=360)
+def infer_gpu_slow(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu):
+    pipe = pipe.to("cuda") 
+    generator = torch.Generator(device='cuda')
+    generator = generator.manual_seed(seed)
+    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)
     result = pipe(prompt, negative_prompt=negative_prompt, generator=generator,
                 num_inference_steps=ddim_steps, guidance_scale=guidance_scale,
                 resolutions_list=resolutions_list, fast_mode=fast_mode, cosine_scale=cosine_scale,
@@ -26,27 +57,28 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
     if output_size == "2048 x 2048":
         resolutions_list = [[1024, 1024],
                             [2048, 2048]]
+        infer_gpu_part = infer_gpu_fast
     elif output_size == "2048 x 4096":
         resolutions_list = [[512, 1024],
                             [1024, 2048],
                             [2048, 4096]]
+        infer_gpu_part = infer_gpu_mid
     elif output_size == "4096 x 2048":
         resolutions_list = [[1024, 512],
                             [2048, 1024],
                             [4096, 2048]]
+        infer_gpu_part = infer_gpu_mid
     elif output_size == "4096 x 4096":
         resolutions_list = [[1024, 1024],
                             [2048, 2048],
                             [4096, 4096]]
+        infer_gpu_part = infer_gpu_slow
 
     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)
 
     print('GPU starts')
-    result = infer_gpu_part(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale)
+    result = infer_gpu_part(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu)
     print('GPU ends')
 
     image = result.images[0]