import gradio as gr
import torch
from PIL import Image
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler

MODEL_ID = "timbrooks/instruct-pix2pix"

pipe = None

def load_pipe():
    global pipe
    if pipe is not None:
        return pipe

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    p = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        safety_checker=None,   # keeps it simple for demo; you can re-add later
    )

    # Scheduler that often looks better for edits
    p.scheduler = EulerAncestralDiscreteScheduler.from_config(p.scheduler.config)

    p = p.to(device)
    pipe = p
    return pipe

def edit_image(image: Image.Image, prompt: str, strength: float, guidance: float, steps: int, seed: int):
    if image is None:
        return None

    if not prompt or not prompt.strip():
        return image

    p = load_pipe()
    device = "cuda" if torch.cuda.is_available() else "cpu"

    generator = torch.Generator(device=device)
    if seed >= 0:
        generator = generator.manual_seed(seed)

    # InstructPix2Pix expects an RGB PIL image
    image = image.convert("RGB")

    out = p(
        prompt=prompt,
        image=image,
        num_inference_steps=int(steps),
        image_guidance_scale=float(strength),   # how much it follows the input image
        guidance_scale=float(guidance),         # how much it follows the prompt
        generator=generator,
    ).images[0]

    return out

with gr.Blocks(title="SNAP AI Editor") as demo:
    gr.Markdown("## SNAP AI Editor\nUpload an image and describe the edit you want.")

    with gr.Row():
        input_img = gr.Image(type="pil", label="Input image")
        output_img = gr.Image(type="pil", label="Output image")

    prompt = gr.Textbox(
        label="Prompt",
        placeholder="Examples: 'put me in a tuxedo', 'remove acne and smooth skin', 'change hair to blonde'"
    )

    with gr.Row():
        strength = gr.Slider(0.5, 2.0, value=1.2, step=0.05, label="Keep Original (image_guidance)")
        guidance = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="Follow Prompt (guidance)")
        steps = gr.Slider(10, 40, value=25, step=1, label="Steps")
        seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")

    btn = gr.Button("Submit")
    btn.click(
        fn=edit_image,
        inputs=[input_img, prompt, strength, guidance, steps, seed],
        outputs=[output_img],
    )

demo.queue().launch()