import gradio as gr
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from PIL import Image

def load_model():
    controlnet = ControlNetModel.from_pretrained(
        "Yuanshi/OminiControl",
        torch_dtype=torch.float16,
        use_safetensors=True
    )
    
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        controlnet=controlnet,
        torch_dtype=torch.float16,
        safety_checker=None
    ).to("cuda" if torch.cuda.is_available() else "cpu")
    return pipe

def generate(image, prompt, resolution):
    pipe = load_model()
    output = pipe(
        prompt=prompt,
        image=image,
        num_inference_steps=20,
        controlnet_conditioning_scale=1.0,
        width=resolution,
        height=resolution
    ).images[0]
    return output

# Create Gradio interface
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Enter your prompt"),
        gr.Radio(choices=[512, 1024], value=512, label="Resolution")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="OminiControl Image Editor",
    description="Upload an image and provide a prompt to edit it."
)

if __name__ == "__main__":
    demo.launch()