import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig, BitsAndBytesConfig
import os
#import subprocess
#subprocess.run("pip list", shell=True)
#subprocess.run("diffusers-cli env", shell=True)
#from optimum.quanto import freeze, qfloat8, quantize

HF_TOKEN = os.getenv("HF_TOKEN", "")
device = "cuda" if torch.cuda.is_available() else "cpu"
flux_repo = "multimodalart/FLUX.1-dev2pro-full"
ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
transformer_gguf = FluxTransformer2DModel.from_single_file(ckpt_path, subfolder="transformer", quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
                                                           torch_dtype=torch.bfloat16, config=flux_repo, token=HF_TOKEN)
transformer = FluxTransformer2DModel.from_pretrained(flux_repo, subfolder="transformer", torch_dtype=torch.bfloat16, token=HF_TOKEN)
nf4_quantization_config = BitsAndBytesConfig(load_in_4bit=True)
transformer_nf4 = FluxTransformer2DModel.from_pretrained(flux_repo, subfolder="transformer", quantization_config=nf4_quantization_config,
                                                         torch_dtype=torch.bfloat16, token=HF_TOKEN)
pipe = FluxPipeline.from_pretrained(flux_repo, transformer=transformer, torch_dtype=torch.bfloat16, token=HF_TOKEN)
hyper_sd_lora = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")

@spaces.GPU(duration=70)
def infer(prompt: str, mode: str, is_lora: bool, progress=gr.Progress(track_tqdm=True)):
    global pipe
    try:
        pipe.unload_lora_weights()
        if mode == "Default": pipe.transformer = transformer
        elif mode == "GGUF": pipe.transformer = transformer_gguf
        elif mode == "NF4": pipe.transformer = transformer_nf4
        if is_lora:
            pipe.load_lora_weights(hyper_sd_lora, adapter_name="hyper-sd")
            pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
            steps = 8
        else: steps = 28
        pipe.to(device)
        image = pipe(prompt, generator=torch.manual_seed(0), num_inference_steps=steps).images[0]
        pipe.to("cpu")
        return image
    except Exception as e:
        raise gr.Error(e)

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value="A cat holding a sign that says hello world", lines=1)
            mode = gr.Radio(label="Mode", choices=["Default", "GGUF", "NF4"], value="Default")
            is_lora = gr.Checkbox(label="Enable LoRA", value=True)
            gen_btn = gr.Button("Generate Image")
        with gr.Column():
            result = gr.Image(label="Result Image")

    gen_btn.click(infer, [prompt, mode, is_lora], [result])

demo.launch()