Create app.py

app.py

@@ -0,0 +1,79 @@
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from PIL import Image, ImageFilter
+import numpy as np
+import gradio as gr
+import cv2
+
+# Load pre-trained Stable Diffusion model (frozen part)
+model_id = "runwayml/stable-diffusion-v1-5"
+controlnet_id = "lllyasviel/control_v11p_sd15_canny"  # ControlNet for edge detection-based control
+
+# Load ControlNet model (trainable part)
+controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
+
+# Load Stable Diffusion pipeline with ControlNet
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    model_id, controlnet=controlnet, torch_dtype=torch.float16
+)
+
+# Use an efficient scheduler
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+
+# Move pipeline to GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe.to(device)
+
+# Function to generate control image (edge detection using Canny filter)
+def generate_control_image(input_image_path):
+    image = cv2.imread(input_image_path, cv2.IMREAD_GRAYSCALE)
+    edges = cv2.Canny(image, 100, 200)  # Apply Canny edge detection
+    control_image = Image.fromarray(edges).convert("L")
+    control_image = control_image.resize((512, 512))  # Resize to match model requirements
+    control_image.save("control_image.jpg")
+    return "control_image.jpg"
+
+# Function to apply color change
+def apply_color_change(input_image, prompt):
+    # Save input image temporarily
+    input_image_path = "input_image.jpg"
+    input_image.save(input_image_path)
+
+    # Generate control image (edges)
+    control_image_path = generate_control_image(input_image_path)
+
+    # Load processed input and control images
+    input_image = Image.open(input_image_path).convert("RGB").resize((512, 512))
+    control_image = Image.open(control_image_path).convert("L")
+
+    # Generate the new image using the pipeline
+    generator = torch.manual_seed(42)  # For reproducibility
+    output_image = pipe(
+        prompt=prompt,
+        image=input_image,
+        control_image=control_image,
+        generator=generator,
+        num_inference_steps=30
+    ).images[0]
+
+    output_image.save("output_color_changed.png")
+    return "output_color_changed.png"
+
+# Gradio interface
+def gradio_interface(input_image, prompt):
+    output_image_path = apply_color_change(input_image, prompt)
+    return output_image_path
+
+# Launch the Gradio interface with drag and drop
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Image(type="pil", label="Upload your image"),  # Drag and drop feature
+        gr.Textbox(label="Enter prompt", placeholder="e.g. A hoodie with blue and white design"),
+    ],
+    outputs=gr.Image(label="Color Changed Output"),
+    title="AI-Powered Clothing Color Changer",
+    description="Upload an image of clothing, enter a prompt, and get a redesigned color version.",
+)
+
+interface.launch()