Update app.py

app.py

@@ -1,49 +1,44 @@
-import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-from diffusers import DiffusionPipeline
 import gradio as gr
+import torch
+import numpy as np
 from PIL import Image
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers.utils import make_image_grid
+import cv2
 
-# Load Stable Diffusion 3 (from InstantX)
-model_id = "stabilityai/stable-diffusion-3-medium"
-
-# Load the ControlNet model (use an appropriate pre-trained controlnet model)
-controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-
-# Set up the pipeline using both SD3 and ControlNet
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16, use_safetensors=True)
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    model_id,
-    controlnet=controlnet,
-    torch_dtype=torch.float16
+    "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
 )
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.enable_model_cpu_offload()
 
-# Use GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe.to(device)
+def generate_image(input_image, text_prompt):
+    original_image = np.array(input_image)
+    low_threshold = 100
+    high_threshold = 200
+    edges = cv2.Canny(original_image, low_threshold, high_threshold)
+    edges = edges[:, :, None]
+    canny_image = np.concatenate([edges, edges, edges], axis=2)
+    canny_image_pil = Image.fromarray(canny_image)
 
-# Function for Img2Img with ControlNet
-def controlnet_img2img(image, prompt, strength=0.8, guidance=7.5):
-    image = Image.fromarray(image).convert("RGB")  # Convert to RGB
-    
-    # Run the pipeline
-    result = pipe(prompt=prompt, image=image, strength=strength, guidance_scale=guidance).images[0]
-    return result
+    output_image = pipe(text_prompt, image=canny_image_pil).images[0]
 
-# Gradio Interface
-def img_editor(input_image, prompt):
-    result = controlnet_img2img(input_image, prompt)
-    return result
+    result_grid = make_image_grid([input_image, canny_image_pil, output_image], rows=1, cols=3)
+    return result_grid
 
-# Create Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Img2Img Editor with ControlNet and Stable Diffusion 3")
-    with gr.Row():
-        image_input = gr.Image(source="upload", type="numpy", label="Input Image")
-        prompt_input = gr.Textbox(label="Prompt")
-    result_output = gr.Image(label="Output Image")
+    gr.Markdown("# Image Transformation with ControlNet and Stable Diffusion")
     
-    submit_btn = gr.Button("Generate")
-    submit_btn.click(fn=img_editor, inputs=[image_input, prompt_input], outputs=result_output)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Image", tool="editor")
+            text_prompt = gr.Textbox(label="Enter a prompt for the transformation")
+
+    generate_button = gr.Button("Generate Image")
+
+    result = gr.Image(label="Result", shape=(768, 256))
+
+    generate_button.click(fn=generate_image, inputs=[input_image, text_prompt], outputs=result)
 
-# Launch Gradio interface
 demo.launch()