import random
import torch
from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline

# Configure deterministic behavior for reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True

MAX_SEED = 2**32 - 1

class ModelManager:
    """
    Handles model initialization, LoRA weight loading, and image generation.
    """
    def __init__(self, base_model: str, lora_repo: str, trigger_word: str = ""):
        self.trigger_word = trigger_word
        self.pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
        self.pipe.load_lora_weights(lora_repo)
        self.pipe.to("cuda")

    def generate_image(self, prompt: str, cfg_scale: float, steps: int, seed: int, 
                       width: int, height: int, lora_scale: float, progress_callback) -> Image.Image:
        """
        Generates an image based on the given prompt and parameters using a callback for progress updates.
        """
        # Establish reproducible generator
        generator = torch.Generator(device="cuda").manual_seed(seed)
        full_prompt = f"{prompt} {self.trigger_word}"

        def callback_fn(step: int, timestep: int, latents):
            percentage = int((step / steps) * 100)
            message = f"Processing step {step} of {steps}..."
            progress_callback(percentage, message)

        # Generate image with integrated progress reporting
        image = self.pipe(
            prompt=full_prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator,
            joint_attention_kwargs={"scale": lora_scale},
            callback=callback_fn,
            callback_steps=1,
        ).images[0]

        return image

# Initialize the model manager with specified models and LoRA weights
model_manager = ModelManager(
    base_model="black-forest-labs/FLUX.1-dev",
    lora_repo="XLabs-AI/flux-RealismLora",
    trigger_word=""
)

def run_generation(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
    """
    Gradio interface callback to manage seed randomization, progress updates, 
    and image generation using the ModelManager.
    """
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Start the progress
    progress(0, "Starting image generation...")

    # Generate the image using the model manager with progress callback integration
    image = model_manager.generate_image(
        prompt, cfg_scale, steps, seed, width, height, lora_scale, progress
    )

    # Mark completion
    progress(100, "Completed!")
    return image, seed

# Example parameters and image path for initializing the interface with defaults
example_image_path = "example0.webp"
example_prompt = (
    "A Jelita Sukawati speaker is captured mid-speech. She has long, dark brown hair that cascades over her shoulders, "
    "framing her radiant, smiling face. Her Latina features are highlighted by warm, sun-kissed skin and bright, "
    "expressive eyes. She gestures with her left hand, displaying a delicate ring on her pinky finger, as she speaks passionately. "
    "The woman is wearing a colorful, patterned dress with a green lanyard featuring multiple badges and logos hanging around her neck. "
    "The lanyard prominently displays the 'CagliostroLab' text. Behind her, there is a blurred background with a white banner "
    "containing logos and text, indicating a professional or conference setting. The overall scene captures the energy and vibrancy "
    "of her presentation."
)
example_cfg_scale = 3.2
example_steps = 32
example_width = 1152
example_height = 896
example_seed = 3981632454
example_lora_scale = 0.85

def load_example():
    # Load example image for initial display
    example_image = Image.open(example_image_path)
    return (
        example_prompt, 
        example_cfg_scale, 
        example_steps, 
        True, 
        example_seed, 
        example_width, 
        example_height, 
        example_lora_scale, 
        example_image
    )

with gr.Blocks() as app:
    gr.Markdown("# Flux RealismLora Image Generator")
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.TextArea(label="Prompt", placeholder="Type a prompt", lines=5)
            generate_button = gr.Button("Generate")
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=example_cfg_scale)
            steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=example_steps)
            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=example_width)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=example_height)
            randomize_seed = gr.Checkbox(True, label="Randomize seed")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=example_seed)
            lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=example_lora_scale)
        with gr.Column(scale=1):
            result = gr.Image(label="Generated Image")
            gr.Markdown(
                "Generate images using RealismLora and a text prompt.\n"
                "[[non-commercial license, Flux.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)]"
            )

    # Load example data on launch
    app.load(
        load_example, 
        inputs=[], 
        outputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, result]
    )
    
    # Set up button interaction
    generate_button.click(
        run_generation,
        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
        outputs=[result, seed]
    )

app.queue()
app.launch()