Update app.py

app.py

@@ -97,21 +97,34 @@ def create_hdr_effect(original_image, hdr):
     hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
     return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
 
+import time
+
 class LazyLoadPipeline:
     def __init__(self):
         self.pipe = None
 
     def load(self):
         if self.pipe is None:
-            self.pipe = self.setup_pipeline()
-            if ENABLE_CPU_OFFLOAD:
-                self.pipe.enable_model_cpu_offload()
-            else:
-                self.pipe.to(device)
-            if USE_TORCH_COMPILE:
-                self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
+            print("Starting to load the pipeline...")
+            start_time = time.time()
+            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)
+                print(f"Pipeline loaded successfully in {time.time() - start_time:.2f} seconds")
+            except Exception as e:
+                print(f"Error loading pipeline: {str(e)}")
+                raise
 
     def setup_pipeline(self):
+        print("Setting up the pipeline...")
         controlnet = ControlNetModel.from_single_file(
             "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
         )
@@ -132,6 +145,7 @@ class LazyLoadPipeline:
         )
         pipe.vae = vae
 
+        print("Loading textual inversions and LoRA weights...")
         pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
         pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
         pipe.load_lora_weights("models/Lora/SDXLrender_v2.0.safetensors")
@@ -150,31 +164,40 @@ lazy_pipe = LazyLoadPipeline()
 
 @spaces.GPU
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
+    print("Starting image processing...")
     torch.cuda.empty_cache()
-    lazy_pipe.load()
-    lazy_pipe.pipe.unet.set_attn_processor(AttnProcessor2_0())
-    
-    condition_image = resize_and_upscale(input_image, resolution)
-    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),
-    }
-    
-    result = lazy_pipe(**options).images[0]
-    return result
+    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)}")
 
 # Gradio interface
 with gr.Blocks() as demo: