import torch
import spaces
import os
from diffusers.utils import load_image
from diffusers.hooks import apply_group_offloading
from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import T5EncoderModel
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from peft import PeftModel, PeftConfig
# from attention_map_diffusers import (
#     attn_maps,
#     init_pipeline,
#     save_attention_maps
# )
import gradio as gr
huggingface_token = os.getenv("HUGGINFACE_TOKEN")
MAX_SEED = 1000000

# quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
# text_encoder_2_8bit = T5EncoderModel.from_pretrained(
#     "LPX55/FLUX.1-merged_uncensored",
#     subfolder="text_encoder_2",
#     quantization_config=quant_config,
#     torch_dtype=torch.bfloat16,
#     token=huggingface_token
# )
text_encoder_2_unquant = T5EncoderModel.from_pretrained(
    "LPX55/FLUX.1-merged_uncensored",
    subfolder="text_encoder_2",
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)
# good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")

# Load pipeline
# controlnet = FluxControlNetModel.from_pretrained(
#     "jasperai/Flux.1-dev-Controlnet-Upscaler",
#     torch_dtype=torch.bfloat16
# )
pipe = FluxControlNetPipeline.from_pretrained(
    "LPX55/FLUX.1M-8step_upscaler-cnet",
    torch_dtype=torch.bfloat16,
    text_encoder_2=text_encoder_2_unquant,
    token=huggingface_token
)
# adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
# adapter_id2 = "XLabs-AI/flux-RealismLora"
# adapter_id3 = "enhanceaiteam/Flux-uncensored-v2"

pipe.to("cuda")

# try:
#     pipe.vae.enable_slicing()
# except:
#     print("debug-2")
# try:
#     pipe.vae.enable_tiling()
# except:
#     print("debug-3")

# pipe.load_lora_weights(adapter_id, adapter_name="turbo")
# pipe.load_lora_weights(adapter_id2, adapter_name="real")
# pipe.load_lora_weights(adapter_id3, weight_name="lora.safetensors", adapter_name="enhance")
# pipe.set_adapters(["turbo", "real", "enhance"], adapter_weights=[0.9, 0.66, 0.6])
# pipe.fuse_lora(adapter_names=["turbo", "real", "enhance"], lora_scale=1.0)
# pipe.unload_lora_weights()
# save to the Hub
# pipe.push_to_hub("FLUX.1M-8step_upscaler-cnet")

@spaces.GPU()
def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
    generator = torch.Generator().manual_seed(seed)

    # Load control image
    control_image = load_image(control_image)
    w, h = control_image.size
    w = w - w % 32
    h = h - h % 32
    control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
    with torch.inference_mode():
        image = pipe(
            generator=generator,
            prompt=prompt,
            control_image=control_image,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            height=control_image.size[1],
            width=control_image.size[0],
            control_guidance_start=0.0,
            control_guidance_end=guidance_end,
        ).images[0]
        
    return image

# Create Gradio interface with rows and columns
with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
    gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
    with gr.Row():
        control_image = gr.Image(type="pil", label="Control Image", show_label=False)
        generated_image = gr.Image(type="pil", label="Generated Image", format="png", show_label=False)
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(lines=4, placeholder="Enter your prompt here...", label="Prompt")
            scale = gr.Slider(1, 3, value=1, label="Scale", step=0.25)
            generate_button = gr.Button("Generate Image", variant="primary")
        with gr.Column(scale=1):
            seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
            steps = gr.Slider(2, 16, value=8, label="Steps")
            controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale")
            guidance_scale = gr.Slider(1, 30, value=3.5, label="Guidance Scale")
            guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End")
            
    
    with gr.Row():
        gr.Markdown("**Tips:** 8 steps is all you need!")
    
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end],
        outputs=[generated_image]
    )

# Launch the app
iface.launch()