import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load the image enhancement model
model_id = "fotographerai/zenctrl_tools"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe.to(device)

def enhance_image(image):
    # Process the image using the AI model
    result = pipe(image).images[0]  # Get the first enhanced image
    return result

# Set up the Gradio interface
interface = gr.Interface(
    fn=enhance_image,
    inputs=gr.Image(type="pil", label="Upload Image to Enhance"),
    outputs=gr.Image(type="pil", label="Enhanced Image"),
    title="AI Image Enhancement",
    description="Enhance your images using FotographerAI's ZenCtrl tools with AI!",
    live=True
)

# Launch the interface
interface.launch()