import gradio as gr
import torch
import numpy as np
import gc
from diffusers import FluxImg2ImgPipeline
from PIL import Image

# 1. SETUP
MODEL_ID = "black-forest-labs/FLUX.1-schnell"
DEVICE = "cpu"
DTYPE = torch.bfloat16

print(f"--- Loading {MODEL_ID} on {DEVICE} ---")
pipe = FluxImg2ImgPipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE)
print("--- Model Loaded ---")

# 2. SYMMETRY LOGIC
def inject_symmetry_lock(image, side="Left"):
    if image is None: return None
    img_array = np.array(image.convert("RGB"))
    height, width, _ = img_array.shape
    midpoint = width // 2
    
    if side == "Left":
        left_side = img_array[:, :midpoint, :]
        right_side = np.fliplr(left_side)
        if right_side.shape[1] != left_side.shape[1]:
             right_side = right_side[:, :left_side.shape[1], :]
        locked_data = np.concatenate((left_side, right_side), axis=1)
    else:
        right_side = img_array[:, midpoint:, :]
        left_side = np.fliplr(right_side)
        locked_data = np.concatenate((left_side, right_side), axis=1)
        
    return Image.fromarray(locked_data)

# 3. GENERATION (SAFE MODE)
def process_image(prompt, image_input, side, strength, seed):
    if image_input is None:
        raise gr.Error("Please upload an image.")
    
    gc.collect()
    
    # Step A: Symmetry
    print(">>> Step 1: Injecting Symmetry")
    processed_image = inject_symmetry_lock(image_input, side)
    
    # Step B: EXTREME RESIZE (512px Max)
    # This is the only way to prevent Timeout on Free CPU
    w, h = processed_image.size
    MAX_SIZE = 512
    scale = MAX_SIZE / max(w, h)
    
    new_w = int((w * scale) // 16 * 16)
    new_h = int((h * scale) // 16 * 16)
    
    print(f">>> Step 2: Resizing to {new_w}x{new_h} to prevent crash...")
    processed_image = processed_image.resize((new_w, new_h))
    
    # Step C: Flux Inference
    print(f">>> Step 3: Running Flux (Prompt: {prompt})")
    generator = torch.Generator(DEVICE).manual_seed(int(seed))
    
    try:
        result = pipe(
            prompt=prompt,
            image=processed_image,
            strength=strength,      
            num_inference_steps=4,  # Turbo Steps
            guidance_scale=0.0,
            generator=generator
        ).images[0]
        
        return result
    except Exception as e:
        print(f"ERROR: {e}")
        return None

# 4. UI
css = """
#col-container { max-width: 900px; margin: 0 auto; background-color: #f0f2f6; padding: 20px; border-radius: 10px; }
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# ⚡ Flux Face Symmetry (Safe Mode)")
        gr.Markdown("Optimized for CPU Stability (Max 512px).")
        
        with gr.Row():
            with gr.Column():
                img_in = gr.Image(label="Upload Face", type="pil")
                side = gr.Radio(["Left", "Right"], label="Best Side", value="Left")
                prompt = gr.Text(label="Prompt", value="perfect symmetry, smooth skin")
                strength = gr.Slider(0.15, 0.45, value=0.25, step=0.01, label="Strength")
                seed = gr.Number(label="Seed", value=123)
                btn = gr.Button("Generate", variant="primary")
            
            with gr.Column():
                img_out = gr.Image(label="Result")

    btn.click(process_image, inputs=[prompt, img_in, side, strength, seed], outputs=[img_out])

# Enable Queue to handle timeouts better
demo.queue(max_size=2).launch()