import gradio as gr
import numpy as np
import random
import spaces
import torch
import os
from diffusers import Flux2KleinPipeline
from huggingface_hub import hf_hub_download, InferenceClient
from safetensors import safe_open
from safetensors.torch import load_file

dtype = torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"
hf_token = os.getenv("HF_TOKEN")

# Pipeline yükle
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    torch_dtype=dtype,
    token=hf_token
).to(device)

# LoRA yükle
LORA_PATH = "safetensors/IlgaCengizFLUX.2-klein-base-9b.safetensors"
pipe.load_lora_weights(LORA_PATH)
pipe.fuse_lora(lora_scale=0.75)

MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU() # ZeroGPU'yu tetikleyen dekoratör
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024,
          num_inference_steps=28, guidance_scale=3.5,
          progress=gr.Progress(track_tqdm=True)):

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # FLUX'ta textual inversion yok; trigger word'ü düz metne ekle
    trigger_word = "IlgaCengiz, "  # ← LoRA'nı nasıl eğittiysen o trigger word
    enhanced_prompt = trigger_word + prompt

    generator = torch.Generator(device="cpu").manual_seed(seed)

    image = pipe(
        prompt=enhanced_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
    ).images[0]

    return image, seed

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# Ilga Character Generator\nFLUX.2-klein + Custom LoRA")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width = gr.Slider(label="Width", minimum=512, maximum=1360, step=64, value=1024)
                height = gr.Slider(label="Height", minimum=512, maximum=1360, step=64, value=1024)

            with gr.Row():
                num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=50, step=1, value=28)
                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.5, value=3.5)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
        outputs=[result, seed]
    )

demo.launch()