import os
import gradio as gr
import numpy as np
import json
from accelerate import dispatch_model, infer_auto_device_map
from accelerate.utils import get_balanced_memory
from torch.cuda.amp import autocast
import torch
import spaces  # Import this first to avoid CUDA initialization issues
import random
import time
from PIL import Image
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast

# Use the 'waffles' environment variable as the access token
hf_token = os.getenv('waffles')

# Ensure the token is loaded correctly
if not hf_token:
    raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.")

# Define the device
dtype = torch.bfloat16
device = "cuda:0" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    device = torch.device("cuda")
    n_gpu = torch.cuda.device_count()
    torch.cuda.get_device_name(0)
else:
    device = torch.device("cpu")

count0 = torch.zeros(1).to(device)
count1 = torch.zeros(1).to(device)
count2 = torch.zeros(1).to(device)

# Load LoRAs from JSON file
with open('loras.json', 'r') as f:
    loras = json.load(f)

# Initialize the base model with authentication and specify the device
# Initialize the base model with authentication and specify the device
pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token).to(device)

MAX_SEED = 2**32 - 1
MAX_IMAGE_SIZE = 2048


class calculateDuration:
    def __init__(self, activity_name=""):
        self.activity_name = activity_name

    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.activity_name:
            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
        else:
            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")

@spaces.GPU(duration=90)
def generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress):
    generator = torch.Generator(device=device).manual_seed(seed)
    images = []
    
    with calculateDuration("Generating images"):
        for _ in range(num_images):
            # Generate each image
            image = pipe(
                prompt=f"{prompt} {trigger_word}",
                num_inference_steps=steps,
                guidance_scale=cfg_scale,
                width=width,
                height=height,
                generator=generator,
                joint_attention_kwargs={"scale": lora_scale},
            ).images[0]
            images.append(image)
    return images

def run_lora(prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, width, height, lora_scale, num_images, progress=gr.Progress(track_tqdm=True)):
    if not selected_repo:
        raise gr.Error("You must select a LoRA before proceeding.")

    selected_lora = next((lora for lora in loras if lora["repo"] == selected_repo), None)
    if not selected_lora:
        raise gr.Error("Selected LoRA not found.")

    lora_path = selected_lora["repo"]
    trigger_word = selected_lora["trigger_word"]

    # Load LoRA weights
    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
        if "weights" in selected_lora:
            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
        else:
            pipe.load_lora_weights(lora_path)
        
    # Set random seed for reproducibility
    with calculateDuration("Randomizing seed"):
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
    
    images = generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress)
    pipe.to("cuda")
    pipe.unload_lora_weights()
    return images, seed  

def update_selection(evt: gr.SelectData):
    index = evt.index
    selected_lora = loras[index]
    return f"Selected LoRA: {selected_lora['title']}", selected_lora["repo"]

run_lora.zerogpu = True

css = '''
#gen_btn{height: 100%}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: auto; width: auto;}
#gallery .gallery-item{width: 50px; height: 50px; margin: 0px;} /* Make buttons 50% height and width */
#gallery img{width: 100%; height: 100%; object-fit: cover;} /* Resize images to fit buttons */
#info_blob {
    background-color: #f0f0f0;
    border: 2px solid #ccc;
    padding: 10px;
    margin: 10px 0;
    text-align: center;
    font-size: 1.2em;
    font-weight: bold;
    color: #333;
    border-radius: 8px;
}
'''
with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
    title = gr.HTML(
        """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> SOONfactory on Schnell LoRas </h1>""",
        elem_id="title",
    )

    # Info blob stating what the app is running
    info_blob = gr.HTML(
        """<div id="info_blob"> Activist, Futurist, and Realist LoRa-stocked Quick-Use Image Manufactory (over Flux Schnell)</div>"""
    )

    selected_lora_text = gr.Markdown("Selected LoRA: None")
    selected_repo = gr.State(value="")

    # Prompt takes the full line
    prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Type a prompt after selecting a LoRA", elem_id="full_line_prompt")

    with gr.Row():
        with gr.Column(scale=1):  # LoRA collection on the left
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="LoRA Gallery",
                allow_preview=False,
                columns=3,
                elem_id="gallery"
            )
        with gr.Column(scale=1):  # Generated images on the right
            result = gr.Gallery(label="Generated Images")
            seed = gr.Number(label="Seed", value=0, interactive=False)

    with gr.Column():
        with gr.Row():
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=1)
            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=4)
        
        with gr.Row():
            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
        
        with gr.Row():
            randomize_seed = gr.Checkbox(True, label="Randomize seed")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.95)
            num_images = gr.Slider(label="Number of Images", minimum=1, maximum=4, step=1, value=1)

    gallery.select(
        fn=update_selection,
        inputs=[],
        outputs=[selected_lora_text, selected_repo]
    )

    generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
    generate_button.click(
        run_lora,
        inputs=[prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, width, height, lora_scale, num_images],
        outputs=[result, seed]
    )

app.queue()
app.launch()