Spaces:
Running
Running
File size: 2,929 Bytes
1e3319b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 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() |