Update app.py

app.py

@@ -11,6 +11,7 @@ import numpy as np
 from diffusers.models.attention_processor import AttnProcessor2_0
 import gradio as gr
 import spaces
+import time
 
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
@@ -52,8 +53,6 @@ def download_models():
 
 download_models()
 
-import time
-
 def timer_func(func):
     def wrapper(*args, **kwargs):
         start_time = time.time()
@@ -71,35 +70,20 @@ class LazyLoadPipeline:
     def load(self):
         if self.pipe is None:
             print("Starting to load the pipeline...")
-            try:
-                self.pipe = self.setup_pipeline()
-                if ENABLE_CPU_OFFLOAD:
-                    print("Enabling CPU offload...")
-                    self.pipe.enable_model_cpu_offload()
-                else:
-                    print(f"Moving pipeline to device: {device}")
-                    self.pipe.to(device)
-                if USE_TORCH_COMPILE:
-                    print("Compiling the model...")
-                    self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
-            except Exception as e:
-                print(f"Error loading pipeline: {str(e)}")
-                raise
+            self.pipe = self.setup_pipeline()
+            print(f"Moving pipeline to device: {device}")
+            self.pipe.to(device)
+            if USE_TORCH_COMPILE:
+                print("Compiling the model...")
+                self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
 
     @timer_func
     def setup_pipeline(self):
         print("Setting up the pipeline...")
-        start_time = time.time()
         controlnet = ControlNetModel.from_single_file(
             "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
         )
-        print(f"ControlNet loaded in {time.time() - start_time:.2f} seconds")
-
-        start_time = time.time()
         safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
-        print(f"Safety checker loaded in {time.time() - start_time:.2f} seconds")
-
-        start_time = time.time()
         model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
         pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
             model_path,
@@ -108,37 +92,21 @@ class LazyLoadPipeline:
             use_safetensors=True,
             safety_checker=safety_checker
         )
-        print(f"Main pipeline loaded in {time.time() - start_time:.2f} seconds")
-
-        start_time = time.time()
         vae = AutoencoderKL.from_single_file(
             "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
             torch_dtype=torch.float16
         )
         pipe.vae = vae
-        print(f"VAE loaded in {time.time() - start_time:.2f} seconds")
-
-        print("Loading textual inversions and LoRA weights...")
-        start_time = time.time()
         pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
         pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
-        print(f"Textual inversions loaded in {time.time() - start_time:.2f} seconds")
-
-        start_time = time.time()
         pipe.load_lora_weights("models/Lora/SDXLrender_v2.0.safetensors")
         pipe.fuse_lora(lora_scale=0.5)
         pipe.load_lora_weights("models/Lora/more_details.safetensors")
-        print(f"LoRA weights loaded in {time.time() - start_time:.2f} seconds")
-
-        start_time = time.time()
         pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
         pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
-        print(f"Scheduler and FreeU set up in {time.time() - start_time:.2f} seconds")
-
         return pipe
 
     def __call__(self, *args, **kwargs):
-        self.load()
         return self.pipe(*args, **kwargs)
 
 class LazyRealESRGAN:
@@ -173,7 +141,7 @@ def resize_and_upscale(input_image, resolution):
     else:
         img = lazy_realesrgan_x4.predict(img)
     return img
-    
+
 @timer_func
 def create_hdr_effect(original_image, hdr):
     if hdr == 0:
@@ -189,44 +157,41 @@ def create_hdr_effect(original_image, hdr):
     return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
 
 lazy_pipe = LazyLoadPipeline()
+lazy_pipe.load()  # Load the pipeline outside of the GPU function
+
+def prepare_image(input_image, resolution, hdr):
+    condition_image = resize_and_upscale(input_image, resolution)
+    condition_image = create_hdr_effect(condition_image, hdr)
+    return condition_image
 
 @spaces.GPU
 @timer_func
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
     print("Starting image processing...")
     torch.cuda.empty_cache()
-    try:
-        lazy_pipe.load()
-        lazy_pipe.pipe.unet.set_attn_processor(AttnProcessor2_0())
-        
-        print("Resizing and upscaling image...")
-        condition_image = resize_and_upscale(input_image, resolution)
-        print("Applying HDR effect...")
-        condition_image = create_hdr_effect(condition_image, hdr)
-        
-        prompt = "masterpiece, best quality, highres"
-        negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
-        
-        options = {
-            "prompt": prompt,
-            "negative_prompt": negative_prompt,
-            "image": condition_image,
-            "control_image": condition_image,
-            "width": condition_image.size[0],
-            "height": condition_image.size[1],
-            "strength": strength,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": guidance_scale,
-            "generator": torch.Generator(device=device).manual_seed(0),
-        }
-        
-        print("Running inference...")
-        result = lazy_pipe(**options).images[0]
-        print("Image processing completed successfully")
-        return result
-    except Exception as e:
-        print(f"Error during image processing: {str(e)}")
-        raise gr.Error(f"An error occurred: {str(e)}")
+    
+    condition_image = prepare_image(input_image, resolution, hdr)
+    
+    prompt = "masterpiece, best quality, highres"
+    negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
+    
+    options = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "image": condition_image,
+        "control_image": condition_image,
+        "width": condition_image.size[0],
+        "height": condition_image.size[1],
+        "strength": strength,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": torch.Generator(device=device).manual_seed(0),
+    }
+    
+    print("Running inference...")
+    result = lazy_pipe(**options).images[0]
+    print("Image processing completed successfully")
+    return result
 
 # Gradio interface
 with gr.Blocks() as demo: