import spaces 
import torch  
import gradio as gr
from diffusers import DiffusionPipeline
from huggingface_hub import HfApi
import os

# 1. Setup API to automatically fetch LoRAs from your account
api = HfApi()
HF_USERNAME = "sanetium" # Your Hugging Face username

def get_my_loras():
    try:
        # Fetches models from your account. The token is automatically 
        # picked up if you add HF_TOKEN to your Space secrets.
        models = api.list_models(author=HF_USERNAME)
        return ["None"] + [m.id for m in models]
    except Exception as e:
        return ["None"]

# 2. Load the unified FLUX Klein 9B KV model
# Loaded in bfloat16 into system RAM. ZeroGPU handles moving it to VRAM dynamically.
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9b-kv",
    torch_dtype=torch.bfloat16
)

current_lora = None

# 3. ZeroGPU Decorator for serverless inference
@spaces.GPU
def generate_image(prompt, input_image, lora_choice):
    global current_lora
    
    # Dynamically manage LoRA weights to keep features consistent
    if lora_choice != "None" and lora_choice != current_lora:
        if current_lora is not None:
            pipe.unload_lora_weights()
        pipe.load_lora_weights(lora_choice)
        current_lora = lora_choice
    elif lora_choice == "None" and current_lora is not None:
        pipe.unload_lora_weights()
        current_lora = None

    # Distilled KV variant optimal settings for maximum speed
    kwargs = {
        "prompt": prompt,
        "num_inference_steps": 4, 
        "guidance_scale": 1.0,
        "height": 1024,
        "width": 1024,
    }
    
    # Unified generation: if an image is provided, it performs editing
    if input_image is not None:
        kwargs["image"] = input_image
        
    image = pipe(**kwargs).images[0]
    return image

# 4. Interface Construction
# Utilizing a custom Soft theme structure for higher quality UI
theme = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="red"
)

# FIX 1: Removed theme=theme from Blocks
with gr.Blocks() as demo:
    gr.Markdown("# FLUX Klein 9B KV Studio")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="e.g., Change the model's outfit to a red silk summer dress, keeping the pose and face consistent..."
            )
            
            with gr.Row():
                lora_dropdown = gr.Dropdown(
                    choices=get_my_loras(), 
                    value="None", 
                    label="Select LoRA Adapter",
                    info="Dynamically loaded from your HF models"
                )
                refresh_btn = gr.Button("🔄 Refresh", scale=0)
            
            # FIX 2: Merged the help text into the label
            input_img = gr.Image(
                type="pil", 
                label="Reference Image (Optional - Leave empty to generate, upload model photo to edit)"
            )
            
            generate_btn = gr.Button("Generate", variant="primary")
            
        with gr.Column():
            output_img = gr.Image(label="Result")
            
    # Wire up the refresh button to pull new LoRAs without restarting the Space
    refresh_btn.click(
        fn=lambda: gr.update(choices=get_my_loras()), 
        outputs=lora_dropdown
    )
            
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, input_img, lora_dropdown],
        outputs=[output_img]
    )

# FIX 3: Moved theme=theme to the launch command
demo.launch(theme=theme)