import gradio as gr
import torch
from diffusers import FluxImg2ImgPipeline
from PIL import Image
import os

# ---------------------------------------------------------------------------
# FLUX.1 Kontext [dev]
# ⚠️ CRITICAL WARNING: This model is 12B parameters (~24GB).
# 1. This model WILL NOT RUN on a free Hugging Face CPU space (16GB RAM limit).
# 2. It requires a paid GPU instance (A10G, L4, or A100).
# 3. This model is GATED. You must accept the license on Hugging Face 
#    and add your HF_TOKEN as a Secret in your Space settings.
# ---------------------------------------------------------------------------

HF_TOKEN = os.getenv("HF_TOKEN")

print("Attempting to load FLUX.1 Kontext [dev]...")

try:
    # We use bfloat16 for memory efficiency, but this requires a GPU.
    # On CPU, we must use float32, but it will almost certainly OOM.
    pipe = FluxImg2ImgPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-Kontext-dev", 
        torch_dtype=torch.float32, 
        use_auth_token=HF_TOKEN
    )
    # pipe.to("cuda") # Uncomment if using a GPU Space
except Exception as e:
    print(f"FAILED TO LOAD MODEL: {e}")
    pipe = None

def process_image(init_image, prompt, strength, steps):
    if pipe is None:
        return Image.new("RGB", (512, 512), (50, 0, 0)) # Error indicator
        
    init_image = init_image.convert("RGB").resize((512, 512))
    
    # Generate
    image = pipe(
        prompt=prompt, 
        image=init_image, 
        num_inference_steps=int(steps), 
        strength=float(strength),
        guidance_scale=3.5
    ).images[0]
    
    return image

with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# 🪄 WiggleAgent // FLUX Kontext SOTA")
    gr.Markdown("Using FLUX.1-Kontext-dev for high-fidelity in-context editing.")
    
    if not HF_TOKEN:
        gr.Markdown("## ⚠️ ERROR: HF_TOKEN Secret not found in Space Settings!")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Screen")
            prompt = gr.Textbox(label="Edit Prompt", value="redesign the UI with a cyberpunk aesthetic")
            strength = gr.Slider(minimum=0.1, maximum=1.0, value=0.6, label="Edit Strength")
            steps = gr.Slider(minimum=10, maximum=30, value=20, label="Steps")
            btn = gr.Button("Transform", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(type="pil", label="Result")
            
    btn.click(
        fn=process_image, 
        inputs=[input_image, prompt, strength, steps], 
        outputs=output_image,
        api_name="predict"
    )

demo.launch()