import torch
import gradio as gr
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from transformers import CLIPImageProcessor
from PIL import Image
import numpy as np

def load_controlnet_model():
    """Load Stable Diffusion ControlNet pipeline with IP-Adapter."""
    controlnet = ControlNetModel.from_pretrained(
        "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
    )
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        controlnet=controlnet,
        torch_dtype=torch.float16
    ).to("cuda")
    return pipe

def preprocess_image(image):
    """Convert image to depth map for ControlNet input."""
    image = image.convert("L")  # Convert to grayscale
    image = np.array(image)
    depth_map = np.clip(image, 0, 255)
    return Image.fromarray(depth_map)

def generate_image(prompt, input_image):
    """Generate an image using Stable Diffusion ControlNet with IP-Adapter."""
    pipe = load_controlnet_model()
    processed_image = preprocess_image(input_image)
    result = pipe(prompt, image=processed_image, num_inference_steps=50).images[0]
    return result

# Gradio Interface
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Image(type="pil", label="Upload reference image")
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Stable Diffusion with ControlNet and IP-Adapter",
    description="Generate images with precise object placement and consistent style using ControlNet and IP-Adapter.",
)

demo.launch(share=True)