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

# 1. Model Setup
model_id = "timbrooks/instruct-pix2pix"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading model on {device}...")

# Load Pipeline
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float16 if device == "cuda" else torch.float32, 
    safety_checker=None
)
pipe.to(device)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

# 2. Core Function
def edit_image(input_image, instruction, steps, guidance_scale, image_guidance_scale):
    if input_image is None:
        return None
    
    # Resize image for speed optimization on free tier
    width, height = input_image.size
    max_dim = 512
    if width > max_dim or height > max_dim:
        ratio = min(max_dim/width, max_dim/height)
        new_size = (int(width*ratio), int(height*ratio))
        input_image = input_image.resize(new_size, Image.LANCZOS)

    # Run Inference
    images = pipe(
        prompt=instruction, 
        image=input_image, 
        num_inference_steps=steps, 
        guidance_scale=guidance_scale, 
        image_guidance_scale=image_guidance_scale
    ).images
    
    return images[0]

# 3. Create Gradio Interface
with gr.Blocks(theme='gradio/soft') as demo:
    
    gr.HTML("""
        <div style="padding-bottom: 10px;">
            <a href="https://cyc1e5phere.netlify.app" target="_self" style="text-decoration: none; color: #2e7d32; font-weight: bold; font-size: 1.1em; display: flex; align-items: center; transition: opacity 0.3s;">
                <span style="margin-right: 8px; font-size: 1.5em;">&#8592;</span> Back to Home
            </a>
        </div>
    """)

    gr.Markdown("# 🎨 CycleSphere")
    gr.Markdown("Upload an image and enter an instruction to edit it (e.g., 'Turn the sky red' or 'Make it look like a painting').")
    
    with gr.Row():
        with gr.Column():
            original_image = gr.Image(label="Upload Original Image", type="pil")
            instruction_text = gr.Textbox(label="Edit Instruction", placeholder="e.g., Turn the apples into oranges")
            
            with gr.Accordion("Advanced Settings", open=False):
                steps_slider = gr.Slider(minimum=10, maximum=50, value=20, step=1, label="Inference Steps")
                text_guidance = gr.Slider(minimum=1, maximum=20, value=7.5, label="Text Guidance Scale")
                img_guidance = gr.Slider(minimum=1, maximum=5, value=1.5, label="Image Guidance Scale")
            
            run_btn = gr.Button("Start Editing", variant="primary")
            
        with gr.Column():
            result_image = gr.Image(label="Edited Result")

    run_btn.click(
        fn=edit_image, 
        inputs=[original_image, instruction_text, steps_slider, text_guidance, img_guidance], 
        outputs=result_image
    )

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