import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import numpy as np
import cv2
from PIL import Image
import gradio as gr

# Load ControlNet and Stable Diffusion models
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny",
    torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")  # Use GPU in Spaces
pipe.enable_model_cpu_offload()  # Optimize memory usage

# Function to process image with ControlNet
def generate_image(input_image, prompt):
    # Convert input image to numpy array and extract Canny edges
    image = np.array(input_image)
    low_threshold = 100
    high_threshold = 200
    canny_image = cv2.Canny(image, low_threshold, high_threshold)
    canny_image = canny_image[:, :, None]
    canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
    canny_image = Image.fromarray(canny_image)

    # Generate image with ControlNet
    output = pipe(
        prompt,
        image=canny_image,
        num_inference_steps=20,
        controlnet_conditioning_scale=0.8
    ).images[0]
    return output

# Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.Textbox(label="Prompt", placeholder="e.g., 'a futuristic city'")
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Stable Diffusion with ControlNet (Canny)",
    description="Upload an image and enter a prompt to generate a new image guided by Canny edges."
)

# Launch the app
interface.launch()