load in advance

app.py

@@ -1,14 +1,21 @@
 import gradio as gr
+
 import spaces
 
-import os
 import torch
-from PIL import Image
 
 from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 
+from pipeline_freescale import StableDiffusionXLPipeline
+from pipeline_freescale_turbo import StableDiffusionXLPipeline_Turbo
+model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
+model_ckpt_turbo = "stabilityai/sdxl-turbo"
+
+pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=torch.float16)
+pipe_turbo = StableDiffusionXLPipeline_Turbo.from_pretrained(model_ckpt_turbo, torch_dtype=torch.float16)
+
 @spaces.GPU(duration=120)
-def infer_gpu_normal(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
+def infer_gpu_part(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
     pipe = pipe.to("cuda") 
     generator = torch.Generator(device='cuda')
     generator = generator.manual_seed(seed)
@@ -22,8 +29,8 @@ def infer_gpu_normal(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_s
                 ).images[0]
     return result
 
-@spaces.GPU(duration=30)
-def infer_gpu_turbo(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
+@spaces.GPU(duration=40)
+def infer_gpu_part_turbo(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps):
     pipe = pipe.to("cuda") 
     generator = torch.Generator(device='cuda')
     generator = generator.manual_seed(seed)
@@ -39,12 +46,13 @@ def infer_gpu_turbo(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_sc
 
 def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, options, negative_prompt):
 
+    print(prompt)
+    print(negative_prompt)
+
     disable_turbo = 'Disable Turbo' in options
     disable_freeu = 'Disable FreeU' in options
 
     if disable_turbo:
-        from pipeline_freescale import StableDiffusionXLPipeline
-        model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
         fast_mode = True
         if output_size == "2048 x 2048":
             resolutions_list = [[1024, 1024],
@@ -55,12 +63,13 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
         elif output_size == "2048 x 1024":
             resolutions_list = [[1024, 512],
                                 [2048, 1024]]
-        infer_gpu_part = infer_gpu_normal
         restart_steps = [int(ddim_steps * 0.3)] * len(resolutions_list)
 
+        # print('GPU starts')
+        result = infer_gpu_part_turbo(pipe_turbo, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps)
+        # print('GPU ends')
+
     else:
-        from pipeline_freescale_turbo import StableDiffusionXLPipeline
-        model_ckpt = "stabilityai/sdxl-turbo"
         fast_mode = False
         if output_size == "2048 x 2048":
             resolutions_list = [[512, 512],
@@ -74,19 +83,13 @@ def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, o
             resolutions_list = [[512, 256],
                                 [1024, 512],
                                 [2048, 1024]]
-        infer_gpu_part = infer_gpu_turbo
         restart_steps = [int(ddim_steps * 0.5)] * len(resolutions_list)
 
-    pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=torch.float16)
-
-    print('GPU starts')
-    result = infer_gpu_part(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps)
-    print('GPU ends')
+        # print('GPU starts')
+        result = infer_gpu_part(pipe, seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, disable_freeu, restart_steps)
+        # print('GPU ends')
 
-    save_path = 'output.png'
-    result.save(save_path)
-
-    return save_path
+    return result
 
 
 examples = [

pipeline_freescale_turbo.py

@@ -191,7 +191,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     return noise_cfg
 
 
-class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin):
+class StableDiffusionXLPipeline_Turbo(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin):
     r"""
     Pipeline for text-to-image generation using Stable Diffusion XL.