import subprocess
# subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)

import torch
import spaces
import os
import datetime
import io
import moondream as md
from datasets import load_dataset, Dataset, DatasetDict, Image as HFImage
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 LlavaForConditionalGeneration, TextIteratorStreamer, AutoProcessor
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
# from liger_kernel.transformers import apply_liger_kernel_to_llama
from PIL import Image
from threading import Thread
from typing import Generator
# from peft import PeftModel, PeftConfig
import gradio as gr
from huggingface_hub import CommitScheduler, HfApi, logging
from debug import log_params, scheduler, save_image
logging.set_verbosity_debug()

huggingface_token = os.getenv("HUGGINFACE_TOKEN")
MAX_SEED = 1000000

md_api_key = os.getenv("MD_KEY")
model = md.vl(api_key=md_api_key)

text_encoder_2_unquant = T5EncoderModel.from_pretrained(
    "LPX55/FLUX.1-merged_uncensored",
    subfolder="text_encoder_2",
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)

pipe = FluxControlNetPipeline.from_pretrained(
    "LPX55/FLUX.1M-8step_upscaler-cnet",
    torch_dtype=torch.bfloat16,
    text_encoder_2=text_encoder_2_unquant,
    token=huggingface_token
)
pipe.to("cuda")

# torch._dynamo.config.suppress_errors = True
# For FLUX models, compiling VAE decode can also be beneficial if needed, though UNet is primary.
# pipe.vae.decode = torch.compile(pipe.vae.decode, mode="reduce-overhead", fullgraph=True) # Uncomment if VAE compile helps
# try:
#     pipe.vae.decode = torch.compile(pipe.vae.decode, mode="default")
# except Exception as e:
#     print(f"Compile failed: {e}")

# 2. Memory Efficient Attention (xFormers): Reduces memory usage and improves speed
# Requires xformers library installation. Beneficial even with high VRAM.
try:
    pipe.enable_xformers_memory_efficient_attention()
except Exception as e:
    print(f"XFormers not available, skipping memory efficient attention: {e}")

# 3. Attention Slicing: Recommended for FLUX models and high-resolution images,
# even with ample VRAM, as it can sometimes help with very large tensors.
pipe.enable_attention_slicing()

@spaces.GPU(duration=12)
@torch.no_grad()
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]))
    print(f"PromptLog: {repr(prompt)}")
    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]
        # print("Type: " + str(type(image)))
    return image

def combine_caption_focus(caption, focus):
    if caption is None:
        caption = ""
    if focus is None:
        focus = "highly detailed photo, raw photography."
    return (str(caption) + "\n\n" + str(focus)).strip()

def generate_caption(control_image):
    if control_image is None:
        return None, None
    
    # Generate a detailed caption
    mcaption = model.caption(control_image, length="short")
    detailed_caption = mcaption["caption"]
    print(f"Detailed caption: {detailed_caption}")
    
    return detailed_caption

def generate_focus(control_image, focus_list):
    if control_image is None:
        return None
    if focus_list is None:
        return ""
    # Generate a detailed caption
    focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
    focus_description = focus_query["answer"]
    print(f"Areas of focus: {focus_description}") 

    return focus_description

def process_image(control_image, user_prompt, system_prompt, scale, steps, 
                controlnet_conditioning_scale, guidance_scale, seed, 
                guidance_end, temperature, top_p, max_new_tokens, log_prompt):
    # Initialize with empty caption
    final_prompt = user_prompt.strip()
    # If no user prompt provided, generate a caption first
    if not final_prompt:
        # Generate a detailed caption
        print("Generating caption...")
        mcaption = model.caption(control_image, length="normal")
        detailed_caption = mcaption["caption"]
        final_prompt = detailed_caption
        yield f"Using caption: {final_prompt}", None, final_prompt
    
    # Show the final prompt being used
    yield f"Generating with: {final_prompt}", None, final_prompt
    
    # Generate the image
    try:
        image = generate_image(
            prompt=final_prompt,
            scale=scale,
            steps=steps,
            control_image=control_image,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guidance_scale=guidance_scale,
            seed=seed,
            guidance_end=guidance_end
        )
        
        try:
            debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png"))
            save_image("/tmp/" + str(seed) + "output.png", debug_img)
        except Exception as e:
            print("Error 160: " + str(e))
        log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
        yield f"Completed! Used prompt: {final_prompt}", image, final_prompt
    except Exception as e:
        print("Error: " + str(e))
        yield f"Error: {str(e)}", None, None
    
with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
    gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
    with gr.Row():
        with gr.Accordion():
            control_image = gr.Image(type="pil", label="Control Image", show_label=False)
        with gr.Accordion():
            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, info="Enter your prompt here or wait for auto-generation...", label="Image Description")
            focus = gr.Textbox(label="Area(s) of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.", value="clothing material, textures, ethnicity")
            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.25)
            with gr.Row():
                generate_button = gr.Button("Generate Image", variant="primary")
                caption_button = gr.Button("Generate Caption", variant="secondary")
        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", step=1)
            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():
        with gr.Accordion("Auto-Caption settings", open=False, visible=False):
            system_prompt = gr.Textbox(
                lines=4, 
                value="Write a straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with 'This image is…' or similar phrasing.",
                label="System Prompt for Captioning",
                visible=False  # Changed to visible
            )
            temperature_slider = gr.Slider(
                minimum=0.0, maximum=2.0, value=0.6, step=0.05,
                label="Temperature",
                info="Higher values make the output more random, lower values make it more deterministic.",
                visible=False  # Changed to visible
            )
            top_p_slider = gr.Slider(
                minimum=0.0, maximum=1.0, value=0.9, step=0.01,
                label="Top-p",
                visible=False  # Changed to visible
            )
            max_tokens_slider = gr.Slider(
                minimum=1, maximum=2048, value=368, step=1,
                label="Max New Tokens",
                info="Maximum number of tokens to generate. The model will stop generating if it reaches this limit.",
                visible=False  # Changed to visible
            )
        log_prompt = gr.Checkbox(value=True, label="Log", visible=False)  # Changed to visible
    
    gr.Markdown("**Tips:** 8 steps is all you need! Incredibly powerful tool, usage instructions coming soon.")

    caption_state = gr.State()
    focus_state = gr.State()
    log_state = gr.State()

    generate_button.click(
        fn=process_image,
        inputs=[
            control_image, prompt, system_prompt, scale, steps, 
            controlnet_conditioning_scale, guidance_scale, seed, 
            guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt
        ],
        outputs=[log_state, generated_image, prompt]
    )
    control_image.input(
        generate_caption,
        inputs=[control_image],
        outputs=[caption_state]
    ).then(
        generate_focus,
        inputs=[control_image, focus],
        outputs=[focus_state]
    ).then(
        combine_caption_focus,
        inputs=[caption_state, focus_state],
        outputs=[prompt]
    )
    caption_button.click(
        fn=generate_caption,
        inputs=[control_image],
        outputs=[prompt]
    ).then(
        generate_focus,
        inputs=[control_image, focus],
        outputs=[focus_state]
    ).then(
        combine_caption_focus,
        inputs=[caption_state, focus_state],
        outputs=[prompt]
    )

demo.launch(show_error=True)