Update app.py

app.py

@@ -3,21 +3,30 @@ import torch
 import numpy as np
 import cv2
 from PIL import Image
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers import (
+    StableDiffusionControlNetPipeline,
+    ControlNetModel,
+    UniPCMultistepScheduler
+)
 
-# Load ControlNet model
+# Load ControlNet model (Canny)
 controlnet = ControlNetModel.from_pretrained(
     "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32
 )
 
 # Load Stable Diffusion pipeline with ControlNet
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    torch_dtype=torch.float32
 )
 
 # Set the scheduler
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
+# Move the pipeline to the appropriate device
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+
 def process_and_generate(image, prompt, num_inference_steps, guidance_scale):
     # Convert PIL image to numpy array
     image = np.array(image)
@@ -29,14 +38,14 @@ def process_and_generate(image, prompt, num_inference_steps, guidance_scale):
     canny_image = Image.fromarray(image)
 
     # Generate image using the pipeline
-    generated_image = pipe(
-        prompt,
-        canny_image,
+    result = pipe(
+        prompt=prompt,
+        image=canny_image,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-    ).images[0]
+    )
 
-    return generated_image
+    return result.images[0]
 
 # Define the Gradio interface
 iface = gr.Interface(
@@ -48,9 +57,9 @@ iface = gr.Interface(
         gr.Slider(0.1, 10.0, value=7.5, step=0.1, label="Guidance Scale"),
     ],
     outputs=gr.Image(type="pil", label="Generated Image"),
-    title="Stable Diffusion with ControlNet",
-    description="Generate images using Stable Diffusion conditioned on edge maps detected by Canny.",
+    title="🧠 Stable Diffusion with ControlNet (Canny)",
+    description="Upload an image and enter a prompt. The system uses Canny edge detection to guide Stable Diffusion generation.",
 )
 
-# Launch the interface
-iface.launch()
+# Launch the app
+iface.launch()