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

# Device (Hugging Face free CPU)
device = "cpu"

# Load model
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix"
)
pipe = pipe.to(device)

# Optimization (safe)
pipe.enable_attention_slicing()

# Main function
def modify_image(image, text):
    if image is None:
        return None
    
    # Resize for balance (speed + quality)
    image = image.resize((384, 384))
    
    # Generate image
    result = pipe(
        prompt=text,
        image=image,
        num_inference_steps=5,   # ✅ balanced
        guidance_scale=5         # ✅ balanced
    ).images[0]

    return result

# Gradio UI
interface = gr.Interface(
    fn=modify_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Enter Instruction (e.g. convert into basketball court)")
    ],
    outputs=gr.Image(label="Output Image"),
    title="🔥 Fast & Stable Pix2Pix Image Modifier",
    description="Upload image + text → modified image (optimized for Hugging Face CPU)"
)

# Launch
interface.launch()