import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from diffusers import StableDiffusionInstructPix2PixPipeline
import torch

# Load the model (once)
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix",
    torch_dtype=torch.float32,
    safety_checker=None
).to("cpu")

def stylize(image: Image.Image, quote: str) -> Image.Image:
    # Step 1: Stylize with instruct-pix2pix
    edited_image = pipe(image=image, prompt="neon filter", image_guidance_scale=1.5).images[0]

    # Step 2: Overlay text
    draw = ImageDraw.Draw(edited_image)
    font = ImageFont.load_default()
    width, height = edited_image.size
    draw.text((10, height - 40), quote, font=font, fill="cyan")

    return edited_image

interface = gr.Interface(
    fn=stylize,
    inputs=[
        gr.Image(type="pil", label="Upload Your Photo"),
        gr.Textbox(label="Your Quote")
    ],
    outputs=gr.Image(type="pil", label="Neon Image"),
    title="Neon Image Stylizer",
    description="Upload an image of your business or product and add your quote. We'll stylize it with a neon effect!"
)

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