import gradio as gr
import torch
from diffusers import FluxPipeline
import os
import traceback

HF_TOKEN = os.environ.get("HF_TOKEN")

MODEL_IDS = {
    "FLUX.1 Schnell": "black-forest-labs/FLUX.1-schnell",
    "FLUX.1 DEV": "black-forest-labs/FLUX.1-dev",
    "FLUX.1 Kontext": "black-forest-labs/FLUX.1-kontext"
}

PIPELINES = {}

def get_pipeline(model_name):
    try:
        if model_name not in PIPELINES:
            print(f"Loading pipeline for {model_name} ({MODEL_IDS[model_name]})...")
            pipe = FluxPipeline.from_pretrained(
                MODEL_IDS[model_name],
                token=HF_TOKEN,  # Correct argument for Hugging Face tokens
                torch_dtype=torch.float32  # Use float32 for CPU compatibility
            )
            pipe.enable_model_cpu_offload()  # Offload to CPU if needed
            PIPELINES[model_name] = pipe
            print(f"Pipeline for {model_name} loaded successfully.")
        return PIPELINES[model_name]
    except Exception as e:
        print(f"Error loading pipeline for {model_name}: {e}")
        traceback.print_exc()
        raise RuntimeError(f"Failed to load {model_name}: {e}")

def generate_image(model_name, prompt, height, width, guidance_scale, steps, seed):
    print(f"generate_image called with model: {model_name}, prompt: '{prompt}', height: {height}, width: {width}, guidance_scale: {guidance_scale}, steps: {steps}, seed: {seed}")
    if not prompt or not prompt.strip():
        print("Prompt is empty.")
        return None
    try:
        pipe = get_pipeline(model_name)
        generator = torch.Generator("cpu").manual_seed(int(seed))
        images = pipe(
            prompt,
            height=int(height),
            width=int(width),
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(steps),
            max_sequence_length=512,
            generator=generator
        ).images
        print(f"Image generated successfully for prompt: '{prompt}'")
        return images[0]
    except Exception as e:
        print(f"Error during image generation: {e}")
        traceback.print_exc()
        return None  # Gradio will show a blank image

with gr.Blocks() as demo:
    gr.Markdown("# FLUX Text-to-Image Generator\nSelect a model and enter your prompt.")
    with gr.Row():
        model_selector = gr.Dropdown(
            choices=list(MODEL_IDS.keys()),
            value="FLUX.1 Schnell",
            label="Choose FLUX Model"
        )
    prompt = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate")
    with gr.Row():
        height = gr.Slider(256, 1024, value=1024, step=64, label="Height")
        width = gr.Slider(256, 1024, value=1024, step=64, label="Width")
    with gr.Row():
        guidance_scale = gr.Slider(1, 10, value=3.5, step=0.1, label="Guidance Scale")
        steps = gr.Slider(10, 100, value=50, step=1, label="Steps")
        seed = gr.Slider(0, 10000, value=0, step=1, label="Seed")
    generate_btn = gr.Button("Generate Image")
    output_image = gr.Image(type="pil", label="Generated Image")

    generate_btn.click(
        fn=generate_image,
        inputs=[model_selector, prompt, height, width, guidance_scale, steps, seed],
        outputs=output_image
    )

if __name__ == "__main__":
    print("Starting FLUX Text-to-Image Gradio app...")
    demo.launch()