import spaces
import torch
import gradio as gr
import random
from PIL import Image
from diffusers import Flux2KleinPipeline

# ── Configuration ─────────────────────────────────────────────────────────────
BASE_MODEL = "black-forest-labs/FLUX.2-klein-9B"
LORA_REPO = "paom/texture2albedo-v2"
WEIGHT_NAME = "pytorch_lora_weights.safetensors"

DEFAULT_PROMPT = (
    "Unlit flat-shaded albedo map. Remove all shadows, reflections, highlights, "
    "and specularity. Maintain absolute pixel-per-pixel structural identity, shape, "
    "and spatial alignment with the original image, displaying only raw base color."
)

MAX_SEED = 2**31 - 1

# ── Model load at module scope ───────────────────────────────────────────────
print("Loading FLUX.2-klein-9B base pipeline...")
pipe = Flux2KleinPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
print("Loading LoRA weights...")
pipe.load_lora_weights(
    LORA_REPO,
    weight_name=WEIGHT_NAME,
    adapter_name="albedo",
)
# Model card loads the LoRA adapter and runs it unfused (default adapter
# weight = 1.0); do NOT fuse, to match the documented inference recipe.
print("Pipeline ready.")


# ── Inference ────────────────────────────────────────────────────────────────
@spaces.GPU(duration=90)
def generate_albedo(
    input_image,
    prompt,
    num_inference_steps,
    guidance_scale,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if input_image is None:
        raise gr.Error("Please upload a texture or photo first.")

    if not prompt or not prompt.strip():
        prompt = DEFAULT_PROMPT

    orig_width, orig_height = input_image.size

    # Resize to 1024x1024 for the model
    processed_input = input_image.resize((1024, 1024))

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    # Model card seeds with torch.manual_seed(seed), i.e. a CPU generator.
    # A CUDA generator produces a different noise sequence for the same seed,
    # so match the documented recipe to keep outputs consistent with the card.
    generator = torch.manual_seed(seed)

    with torch.inference_mode():
        output_image = pipe(
            prompt=prompt,
            image=processed_input,
            guidance_scale=guidance_scale,
            num_inference_steps=int(num_inference_steps),
            generator=generator,
        ).images[0]

    # Resize back to original dimensions
    albedo_map = output_image.resize((orig_width, orig_height))

    return albedo_map, seed


# ── UI ───────────────────────────────────────────────────────────────────────
with gr.Blocks(title="Texture to Albedo — FLUX.2 Klein") as demo:
    gr.Markdown(
        """
        # Texture → Albedo Studio
        Extract clean, flat, shadowless **albedo maps** from textures and photos using
        [paom/texture2albedo-v2](https://huggingface.co/paom/texture2albedo-v2) on
        [FLUX.2-klein-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B).
        Perfect for 3D/PBR material pipelines.
        """
    )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            input_img = gr.Image(label="Input Texture / Photo", type="pil")
            prompt_box = gr.Textbox(
                label="Prompt",
                value=DEFAULT_PROMPT,
                lines=3,
                placeholder="Describe what you want the albedo map to look like...",
            )

            with gr.Accordion("Advanced Parameters", open=False):
                inference_steps = gr.Slider(
                    minimum=1, maximum=12, value=4, step=1,
                    label="Inference Steps",
                )
                guidance = gr.Slider(
                    minimum=0.0, maximum=4.0, value=1.0, step=0.1,
                    label="Guidance Scale",
                )
                seed_input = gr.Slider(
                    minimum=0, maximum=MAX_SEED, value=0, step=1,
                    label="Seed",
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize seed", value=True,
                )

            submit_btn = gr.Button("Generate Albedo Map", variant="primary", size="lg")

        with gr.Column(scale=1):
            albedo_out = gr.Image(label="Clean Albedo Map", type="pil")
            used_seed = gr.Number(label="Seed used", precision=0, interactive=False)

    gr.Examples(
        # The model-card example images are before/after composites
        # (left half = original texture, right half = albedo output).
        # Only the left "before" half is fed to the model as the example input.
        examples=[
            ["example_1_left.jpg", DEFAULT_PROMPT, 4, 1.0, 0, True],
            ["example_2_left.jpg", DEFAULT_PROMPT, 4, 1.0, 0, True],
            ["example_3_left.jpg", DEFAULT_PROMPT, 4, 1.0, 0, True],
        ],
        inputs=[input_img, prompt_box, inference_steps, guidance, seed_input, randomize_seed],
        outputs=[albedo_out, used_seed],
        fn=generate_albedo,
        cache_examples=True,
        cache_mode="lazy",
    )

    submit_btn.click(
        fn=generate_albedo,
        inputs=[input_img, prompt_box, inference_steps, guidance, seed_input, randomize_seed],
        outputs=[albedo_out, used_seed],
    )

    prompt_box.submit(
        fn=generate_albedo,
        inputs=[input_img, prompt_box, inference_steps, guidance, seed_input, randomize_seed],
        outputs=[albedo_out, used_seed],
    )


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), show_error=True)