init

app.py

@@ -9,35 +9,7 @@ from pipeline_freescale import StableDiffusionXLPipeline
 from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 
 @spaces.GPU(duration=120)
-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):
+def infer_gpu_part(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)
@@ -53,26 +25,16 @@ def infer_gpu_slow(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_sca
 def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, options, negative_prompt):
 
     disable_freeu = 'Disable FreeU' in options
-    fast_mode = 'Fast Mode' in options
+    fast_mode = True
     if output_size == "2048 x 2048":
         resolutions_list = [[1024, 1024],
                             [2048, 2048]]
-        infer_gpu_part = infer_gpu_fast
-    elif output_size == "2048 x 4096":
+    elif output_size == "1024 x 2048":
         resolutions_list = [[512, 1024],
-                            [1024, 2048],
-                            [2048, 4096]]
-        infer_gpu_part = infer_gpu_mid
-    elif output_size == "4096 x 2048":
+                            [1024, 2048]]
+    elif output_size == "2048 x 1024":
         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
+                            [2048, 1024]]
 
     model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
     pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=torch.float16)
@@ -198,7 +160,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Accordion('FreeScale Parameters (feel free to adjust these parameters based on your prompt): ', open=False):
                 with gr.Row():
-                    output_size = gr.Dropdown(["2048 x 2048", "2048 x 4096", "4096 x 2048", "4096 x 4096"], value="2048 x 2048", label="Output Size (H x W)")
+                    output_size = gr.Dropdown(["2048 x 2048", "1024 x 2048", "2048 x 1024"], value="2048 x 2048", label="Output Size (H x W)", info="Due to GPU constraints, run the demo locally for higher resolutions.")
                 with gr.Row():
                     ddim_steps = gr.Slider(label='DDIM Steps',
                              minimum=5,
@@ -222,7 +184,7 @@ with gr.Blocks(css=css) as demo:
                              step=1,
                              value=123)
                 with gr.Row():
-                    options = gr.CheckboxGroup(['Disable FreeU', 'Fast Mode'], label='Options (NOT recommended to change)')
+                    options = gr.CheckboxGroup(['Disable FreeU'], label='Options (NOT recommended to change)')
                 with gr.Row():
                     negative_prompt = gr.Textbox(label='Negative Prompt', value='blurry, ugly, duplicate, poorly drawn, deformed, mosaic')