Update app.py

app.py

@@ -45,10 +45,13 @@ pipe.load_lora_weights(
     weight_name="40kHalf.safetensors"
 )
 
+# Move to CPU initially
+pipe = pipe.to("cpu")
+
 def get_random_condition_image():
     conditions_dir = Path("conditions")
     if conditions_dir.exists():
-        image_files = list(conditions_dir.glob("*.[jp][pn][g]"))  # matches .jpg, .png, .jpeg
+        image_files = list(conditions_dir.glob("*.[jp][pn][g]"))
         if image_files:
             random_image = random.choice(image_files)
             return str(random_image)
@@ -57,38 +60,65 @@ def get_random_condition_image():
 def get_canny_image(image, low_threshold=100, high_threshold=200):
     if isinstance(image, Image.Image):
         image = np.array(image)
+    elif isinstance(image, str):
+        image = np.array(Image.open(image))
     
-    if image.shape[2] == 4:
+    if len(image.shape) == 2:  # If grayscale
+        image = np.stack([image] * 3, axis=-1)
+    elif image.shape[2] == 4:  # If RGBA
         image = image[..., :3]
     
     canny_image = cv2.Canny(image, low_threshold, high_threshold)
     canny_image = np.stack([canny_image] * 3, axis=-1)
     return Image.fromarray(canny_image)
 
-@spaces.GPU(duration=120)
-def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold, seed):
-    if seed is not None and seed != "":
-        try:
-            generator = torch.Generator().manual_seed(int(seed))
-        except ValueError:
+@spaces.GPU(duration=300)  # Increased duration to 5 minutes
+def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold, seed, progress=gr.Progress()):
+    if input_image is None:
+        raise gr.Error("Please provide an input image!")
+    
+    try:
+        if seed is not None and seed != "":
+            try:
+                generator = torch.Generator().manual_seed(int(seed))
+            except ValueError:
+                generator = torch.Generator()
+        else:
             generator = torch.Generator()
-    else:
-        generator = torch.Generator()
 
-    canny_image = get_canny_image(input_image, low_threshold, high_threshold)
-    
-    with torch.no_grad():
-        image = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            image=canny_image,
-            extra_condition_scale=1.0,
-            generator=generator
-        ).images[0]
+        progress(0.1, desc="Processing input image...")
+        canny_image = get_canny_image(input_image, low_threshold, high_threshold)
+        
+        progress(0.2, desc="Moving model to device...")
+        # Move pipeline to GPU
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        pipe.to(device)
+        
+        progress(0.3, desc="Generating image...")
+        with torch.no_grad():
+            image = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=guidance_scale,
+                image=canny_image,
+                extra_condition_scale=1.0,
+                generator=generator
+            ).images[0]
+        
+        progress(0.9, desc="Moving model back to CPU...")
+        # Move back to CPU to free up GPU memory
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
+        
+        progress(1.0, desc="Done!")
+        return canny_image, image
     
-    return canny_image, image
+    except Exception as e:
+        # Move back to CPU in case of error
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
+        raise gr.Error(f"An error occurred: {str(e)}")
 
 def random_image_click():
     image_path = get_random_condition_image()
@@ -99,7 +129,7 @@ def random_image_click():
 # Example data
 examples = [
     [
-        "conditions/example1.jpg",  # Replace with actual paths
+        "conditions/example1.jpg",
         "a futuristic cyberpunk city",
         "blurry, bad quality",
         7.5,
@@ -109,7 +139,7 @@ examples = [
         42
     ],
     [
-        "conditions/example2.jpg",  # Replace with actual paths
+        "conditions/example2.jpg",
         "a serene mountain landscape",
         "dark, gloomy",
         7.0,
@@ -134,6 +164,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="numpy")
             random_image_btn = gr.Button("Load Random Reference Image")
+            status_text = gr.Textbox(label="Status", value="Ready", interactive=False)
             
             prompt = gr.Textbox(
                 label="Prompt",
@@ -195,4 +226,5 @@ with gr.Blocks() as demo:
         outputs=[canny_output, result]
     )
 
+demo.queue()  # Enable queuing for better handling of concurrent requests
 demo.launch()