import gradio as gr
import json
import logging
import argparse
import torch
import os
from os import path
from PIL import Image
import numpy as np
import spaces
import copy
import random
import time
from typing import Any, Dict, List, Optional, Union
from huggingface_hub import hf_hub_download
from diffusers import DiffusionPipeline, FluxTransformer2DModel, FluxPipeline, AutoencoderTiny
import safetensors.torch
from safetensors.torch import load_file
from custom_pipeline import FluxWithCFGPipeline
from transformers import CLIPModel, CLIPProcessor, CLIPConfig
import gc

cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
os.environ["TRANSFORMERS_CACHE"] = cache_path
os.environ["HF_HUB_CACHE"] = cache_path
os.environ["HF_HOME"] = cache_path

device = "cuda" if torch.cuda.is_available() else "cpu"

torch.backends.cuda.matmul.allow_tf32 = True

dtype = torch.bfloat16
pipe = FluxWithCFGPipeline.from_pretrained(
    "ostris/OpenFLUX.1", torch_dtype=dtype
).to("cuda")
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")

pipe.to("cuda")
clipmodel = 'norm'
if clipmodel == "long":
    model_id = "zer0int/LongCLIP-GmP-ViT-L-14"
    config = CLIPConfig.from_pretrained(model_id)
    maxtokens = 77
if clipmodel == "norm":
    model_id = "zer0int/CLIP-GmP-ViT-L-14"
    config = CLIPConfig.from_pretrained(model_id)
    maxtokens = 77
clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=True).to("cuda")
clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=True, return_tensors="pt", truncation=True)
config.text_config.max_position_embeddings = 77

pipe.tokenizer = clip_processor.tokenizer
pipe.text_encoder = clip_model.text_model
pipe.tokenizer_max_length = maxtokens
pipe.text_encoder.dtype = torch.bfloat16
torch.cuda.empty_cache()

# Load LoRAs from JSON file
with open('loras.json', 'r') as f:
    loras = json.load(f)
    
MAX_SEED = 2**32-1

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")


def update_selection(evt: gr.SelectData, width, height):
    selected_lora = loras[evt.index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
    if "aspect" in selected_lora:
        if selected_lora["aspect"] == "portrait":
            width = 768
            height = 1024
        elif selected_lora["aspect"] == "landscape":
            width = 1024
            height = 768
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        width,
        height,
    )

@spaces.GPU(duration=70)
def generate_image(prompt, negative_prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
    pipe.to("cuda")
    generator = torch.Generator(device="cuda").manual_seed(seed)
    
    with calculateDuration("Generating image"):
        # Generate image
        image = pipe(
            prompt=f"{prompt} {trigger_word}",
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator,
            joint_attention_kwargs={"scale": lora_scale},
        ).images[0]
    return image

def run_lora(prompt, negative_prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
    if negative == "":
        negative = None
    if selected_index is None:
        raise gr.Error("Select a LoRA adapter square before proceeding.")

    selected_lora = loras[selected_index]
    lora_path = selected_lora["repo"]
    trigger_word = selected_lora["trigger_word"]
    if(trigger_word):
        if "trigger_position" in selected_lora:
            if selected_lora["trigger_position"] == "prepend":
                prompt_mash = f"{trigger_word} {prompt}"
            else:
                prompt_mash = f"{prompt} {trigger_word}"
        else:
            prompt_mash = f"{trigger_word} {prompt}"
    else:
        prompt_mash = prompt

    # Load LoRA weights
    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
        if "weights" in selected_lora:
            pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"], adapter_name=selected_lora["repo"], adapter_weights=lora_scale)
            pipe.set_adapters(adapter_names={"fast", selected_lora["repo"]}, adapter_weights=[1.0, lora_scale])
        else:
            pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
            pipe.load_lora_weights(lora_path, adapter_name=selected_lora["repo"], adapter_weights=lora_scale)
            pipe.set_adapters(adapter_names={"fast", selected_lora["repo"]}, adapter_weights=[1.0, lora_scale])
        
    # Set random seed for reproducibility
    with calculateDuration("Randomizing seed"):
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
    
    image = generate_image(prompt, negative_prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
    pipe.to("cpu")
    pipe.unload_lora_weights()
    return image, seed  

run_lora.zerogpu = True
#pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
#pipe.set_adapters("fast")
#pipe.set_adapters(["fast", "toy"], adapter_weights=[0.5, 1.0])
#pipe.fuse_lora(adapter_names=["fast"], lora_scale=1.0)

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: 10vh}
'''
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">OpenFlux LoRAsoon®</h1>""",
        elem_id="title",
    )
    	    # Info blob stating what the app is running
    info_blob = gr.HTML(
        """<div id="info_blob"> SOON®'s LoRa Gallery & Art Manufactory Space (running on Ostris' OpenFLUX.1 model w/ their fast-gen LoRA & zer0int's fine-tuned CLIP-GmP-ViT-L-14*! (*the 'normal' 77 tokens variant)) </div>"""
    )

        # Info blob stating what the app is running
    info_blob = gr.HTML(
        """<div id="info_blob"> *Auto-planting of prompts with a choice LoRA trigger bugs on out in this space over flaws yet unclear. In its stead, I've preposed numbered LoRA-box rows with a matched token cheat-sheet: ungainly and free. So, prephrase your best prompts w/: 1-2. HST style autochrome photo |3. RCA style Communist poster |4. SOTS art |5. HST Austin Osman Spare style |6. Vladimir Mayakovsky |7. Marina Tsvetaeva |8. Tsvetaeva_02.CR2 |9. Anna Akhmatova |10. Osip Mandelshtam |11-12. Alexander Blok |13. Blok_02.CR2 |14. LEN Lenin |15. Leon Trotsky |16. Rosa Fluxemburg |17. HST photo in Peterhof |18-19. HST |20. HST autochrome portrait |21. HST archival photo |22. HST 1980s Perestroika-era Soviet photo |23-30. HST style |31. How2Draw a ____ |32. propaganda poster |33. TOK hybrid photo of____ with cartoon of____ |34. 2004 amateur IMG_1099.CR2 photo |35. unexpected photo of |36. flmft |37. 80s analog yearbook photo |38. TOK portra |39. pficonics |40. retrofuturism |41. wh3r3sw4ld0 |42. amateur photo |43. crisp |44-45. IMG_1099.CR2 photo |46. FilmFotos |47. ff-collage |48. vintage cover </div>"""
    )
    selected_index = gr.State(None)
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Select LoRa/Style & type prompt! What do you want to see?")
    with gr.Row():
        with gr.Column(scale=3):
            negative_prompt = gr.Textbox(label="Negative Prompt", lines=1, placeholder="List unwanted conditions, open-fluxedly!")
        with gr.Column(scale=1, elem_id="gen_column"):
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
    with gr.Row():
        with gr.Column(scale=3):
            selected_info = gr.Markdown("")
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="LoRA Inventory",
                allow_preview=False,
                columns=3,
                elem_id="gallery"
            )
            
        with gr.Column(scale=4):
            result = gr.Image(label="Generated Image")

    with gr.Row():
        with gr.Accordion("Advanced Settings", open=True):
            with gr.Column():
                with gr.Row():
                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3)
                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=6)
                
                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, randomize=True)
                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2.0, step=0.01, value=0.85)

    gallery.select(
        update_selection,
        inputs=[width, height],
        outputs=[prompt, selected_info, selected_index, width, height]
    )

    gr.on(
        triggers=[generate_button.click, prompt.submit],
        fn=run_lora,
        inputs=[prompt, negative_prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
        outputs=[result, seed]
    )

app.queue(default_concurrency_limit=None).launch(show_error=True)
app.launch()