load in advance

app.py

@@ -15,15 +15,17 @@ model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
 model_ckpt_turbo = "stabilityai/sdxl-turbo"
 pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=dtype).to(device)
 pipe_turbo = StableDiffusionXLPipeline_Turbo.from_pretrained(model_ckpt_turbo, torch_dtype=dtype).to(device)
+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)
+register_free_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
+register_free_crossattn_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
 torch.cuda.empty_cache()
 
 @spaces.GPU(duration=120)
-def infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
+def infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps):
     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,
@@ -31,13 +33,11 @@ def infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, re
                 ).images[0]
     return result
 
-@spaces.GPU(duration=40)
-def infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
+@spaces.GPU(duration=30)
+def infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps):
     generator = torch.Generator(device='cuda')
     generator = generator.manual_seed(seed)
-    if not disable_freeu:
-        register_free_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
-        register_free_crossattn_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
+
     result = pipe_turbo(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,
@@ -51,7 +51,6 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
     print(negative_prompt)
 
     disable_turbo = 'Disable Turbo' in options
-    disable_freeu = 'Disable FreeU' in options
 
     if disable_turbo:
         fast_mode = True
@@ -66,9 +65,7 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
                                 [2048, 1024]]
         restart_steps = [int(ddim_steps * 0.3)]
 
-        # print('GPU starts')
-        result = infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps)
-        # print('GPU ends')
+        result = infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps)
 
     else:
         fast_mode = False
@@ -86,9 +83,7 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
                                 [2048, 1024]]
         restart_steps = [int(ddim_steps * 0.5)] * 2
 
-        # print('GPU starts')
-        result = infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps)
-        # print('GPU ends')
+        result = infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps)
 
     return result
 
@@ -220,7 +215,7 @@ with gr.Blocks(css=css) as demo:
             with gr.Accordion('Advanced Settings', open=False):
                 with gr.Row():
                     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.", scale=3)
-                    options = gr.CheckboxGroup(['Disable Turbo', 'Disable FreeU'], label="Options", info="NOT recommended to change", scale=2)
+                    options = gr.CheckboxGroup(['Disable Turbo'], label="Options", info="NOT recommended to change", scale=2)
                 with gr.Row():
                     ddim_steps = gr.Slider(label='DDIM Steps',
                              minimum=2,