import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
import logging
from diffusers import DiffusionPipeline, AutoencoderTiny
# Using AttnProcessor2_0 for potential speedup with PyTorch 2.x
from diffusers.models.attention_processor import AttnProcessor2_0
# Assuming custom_pipeline defines FluxWithCFGPipeline correctly
from custom_pipeline import FluxWithCFGPipeline

# --- Setup Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- Torch Optimizations ---
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True # Enable cuDNN benchmark for potentially faster convolutions

# --- Constants ---
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048 # Keep a reasonable limit to prevent OOMs
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_INFERENCE_STEPS = 1 # FLUX Schnell is designed for few steps
MIN_INFERENCE_STEPS = 1
MAX_INFERENCE_STEPS = 8 # Allow slightly more steps for potential quality boost
ENHANCE_STEPS = 2 # Fixed steps for the enhance button

# --- Device and Model Setup ---
dtype = torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = None # Initialize pipe to None

try:
    logging.info("Loading diffusion pipeline...")
    pipe = FluxWithCFGPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
    )
    logging.info("Loading VAE...")
    pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
    
    logging.info(f"Moving pipeline to {device}...")
    pipe.to(device)

    # Apply optimizations
    logging.info("Setting attention processor...")
    pipe.unet.set_attn_processor(AttnProcessor2_0())
    pipe.vae.set_attn_processor(AttnProcessor2_0()) # VAE might benefit too

    logging.info("Loading and fusing LoRA...")
    pipe.load_lora_weights('hugovntr/flux-schnell-realism', weight_name='schnell-realism_v2.3.safetensors', adapter_name="better")
    pipe.set_adapters(["better"], adapter_weights=[1.0])
    pipe.fuse_lora(adapter_name=["better"], lora_scale=1.0) # Fuse for potential speedup
    pipe.unload_lora_weights() # Unload after fusing
    logging.info("LoRA fused and unloaded.")

    # --- Compilation (Major Speed Optimization) ---
    # logging.info("Compiling VAE Decoder...")
    # pipe.vae.decoder = torch.compile(pipe.vae.decoder, mode="reduce-overhead", fullgraph=True)
    # logging.info("Compiling VAE Encoder...")
    # pipe.vae.encoder = torch.compile(pipe.vae.encoder, mode="reduce-overhead", fullgraph=True)
    # logging.info("Model compilation finished.")

    # Clear cache after setup
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        logging.info("CUDA cache cleared after setup.")

except Exception as e:
    logging.error(f"Error during model loading or setup: {e}", exc_info=True)
    # Display error in Gradio if UI is already built, otherwise just log and exit.
    # For simplicity here, we'll rely on the Gradio UI showing an error if `pipe` is None later.
    # If running script directly, consider `sys.exit()`
    # raise gr.Error(f"Failed to load models. Check logs for details. Error: {e}")


# --- Inference Function ---
@spaces.GPU(duration=30) # Slightly increased duration buffer
def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = DEFAULT_INFERENCE_STEPS, is_enhance: bool = False):
    """Generates an image using the FLUX pipeline with error handling."""
    
    if pipe is None:
        raise gr.Error("Diffusion pipeline failed to load. Cannot generate images.")
        
    if not prompt or prompt.strip() == "":
         # Return a blank image or previous result if prompt is empty?
         # For now, raise warning and return None.
         gr.Warning("Prompt is empty. Please enter a description.")
         # Returning None for image, original seed, and error message
         return None, seed, "Error: Empty prompt"

    start_time = time.time()
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Clamp dimensions to avoid excessive memory usage
    width = min(width, MAX_IMAGE_SIZE)
    height = min(height, MAX_IMAGE_SIZE)
    
    # Use fixed steps for enhance button, otherwise use slider value
    steps_to_use = ENHANCE_STEPS if is_enhance else num_inference_steps
    # Clamp steps
    steps_to_use = max(MIN_INFERENCE_STEPS, min(steps_to_use, MAX_INFERENCE_STEPS))

    logging.info(f"Generating image with prompt: '{prompt}', seed: {seed}, size: {width}x{height}, steps: {steps_to_use}")

    try:
        # Ensure generator is on the correct device
        generator = torch.Generator(device=device).manual_seed(int(float(seed)))

        # Use inference_mode for efficiency
        with torch.inference_mode():
            # Generate the image (assuming pipe returns list/tuple with image first)
            # Modify pipe call based on its actual signature if needed
            result_img = pipe(
                prompt=prompt,
                width=width,
                height=height,
                num_inference_steps=steps_to_use,
                generator=generator,
                output_type="pil", # Ensure PIL output for Gradio Image component
                return_dict=False # Assuming the custom pipeline supports this for direct output
            )[0][0] # Assuming the output structure is [[img]]

        latency = time.time() - start_time
        latency_str = f"Latency: {latency:.2f} seconds (Steps: {steps_to_use})"
        logging.info(f"Image generated successfully. {latency_str}")
        return result_img, seed, latency_str

    except torch.cuda.OutOfMemoryError as e:
        logging.error(f"CUDA OutOfMemoryError: {e}", exc_info=True)
        # Clear cache and suggest reducing size/steps
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise gr.Error("GPU ran out of memory. Try reducing the image width/height or the number of inference steps.")
        
    except Exception as e:
        logging.error(f"Error during image generation: {e}", exc_info=True)
        # Clear cache just in case
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise gr.Error(f"An error occurred during generation: {e}")


# --- Real-time Generation Wrapper ---
# This function checks the realtime toggle before calling the main generation function.
# It's triggered by changes in prompt or sliders when realtime is enabled.
def handle_realtime_update(realtime_enabled: bool, prompt: str, seed: int, width: int, height: int, randomize_seed: bool, num_inference_steps: int):
    if realtime_enabled and pipe is not None:
        logging.debug("Realtime update triggered.")
        # Call generate_image directly. Errors within generate_image will be caught and raised as gr.Error.
        # We don't set is_enhance=True for realtime updates.
        return generate_image(prompt, seed, width, height, randomize_seed, num_inference_steps, is_enhance=False)
    else:
        # If realtime is disabled or pipe failed, don't update the image, seed, or latency.
        # Return gr.update() for each output component to indicate no change.
        logging.debug("Realtime update skipped (disabled or pipe error).")
        return gr.update(), gr.update(), gr.update()


# --- Example Prompts ---
examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cute white cat holding a sign that says hello world",
    "an anime illustration of Steve Jobs",
    "Create image of Modern house in minecraft style",
    "photo of a woman on the beach, shot from above. She is facing the sea, while wearing a white dress. She has long blonde hair",
    "Selfie photo of a wizard with long beard and purple robes, he is apparently in the middle of Tokyo. Probably taken from a phone.",
    "Photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.",
    "High-resolution photorealistic render of a sleek, futuristic motorcycle parked on a neon-lit street at night, rain reflecting the lights.",
    "Watercolor painting of a cozy bookstore interior with overflowing shelves and a cat sleeping in a sunbeam.",
]

# --- Gradio UI ---
with gr.Blocks() as demo:
    with gr.Column(elem_id="app-container"):
        gr.Markdown("# 🎨 Realtime FLUX Image Generator")
        gr.Markdown("Generate stunning images in real-time with Modified Flux.Schnell pipeline.")
        gr.Markdown("<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>")

        with gr.Row():
            with gr.Column(scale=2.5):
                result = gr.Image(label="Generated Image", show_label=False, interactive=False)
            with gr.Column(scale=1):
                prompt = gr.Text(
                    label="Prompt",
                    placeholder="Describe the image you want to generate...",
                    lines=3,
                    show_label=False,
                    container=False,
                )
                generateBtn = gr.Button("🖼️ Generate Image")
                enhanceBtn = gr.Button("🚀 Enhance Image")

                with gr.Column("Advanced Options"):
                    with gr.Row():
                        realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=False)
                        latency = gr.Text(label="Latency")
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=42)
                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                    with gr.Row():
                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=MIN_INFERENCE_STEPS, maximum=MAX_INFERENCE_STEPS, step=1, value=DEFAULT_INFERENCE_STEPS)

        with gr.Row():
            gr.Markdown("### 🌟 Inspiration Gallery")
        with gr.Row():
            gr.Examples(
                examples=examples,
                fn=generate_image,
                inputs=[prompt],
                outputs=[result, seed, latency],
                cache_examples=True,
                cache_mode="eager"
            )

    enhanceBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height],
        outputs=[result, seed, latency],
        show_progress="full",
        queue=False,
        concurrency_limit=None
    )

    generateBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="full",
        api_name="RealtimeFlux",
        queue=False
    )

    def update_ui(realtime_enabled):
        return {
            prompt: gr.update(interactive=True),
            generateBtn: gr.update(visible=not realtime_enabled)
        }

    realtime.change(
        fn=update_ui,
        inputs=[realtime],
        outputs=[prompt, generateBtn],
        queue=False,
        concurrency_limit=None
    )

    def realtime_generation(*args):
        if args[0]:  # If realtime is enabled
            return next(generate_image(*args[1:]))

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="full",
        queue=False,
        concurrency_limit=None
    )

    for component in [prompt, width, height, num_inference_steps]:
        component.input(
            fn=realtime_generation,
            inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
            outputs=[result, seed, latency],
            show_progress="hidden",
            trigger_mode="always_last",
            queue=False,
            concurrency_limit=None
        )

# Launch the app
demo.launch()