import os

# Neutralize torch.compile decorators inside mmdit.py (not supported on ZeroGPU
# forked workers) and reduce allocator fragmentation for the 13B DiT.
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import random

import spaces  # noqa: E402  MUST come before torch / CUDA-touching imports
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402

from pipeline import DepthLoRAPipeline  # noqa: E402

MAX_SEED = 2**31 - 1

# --------------------------------------------------------------------- loading
# Turbo base (8-step, no CFG) — the author's recommended fast configuration.
# Krea-2 base checkpoint (~26GB) + depth-control LoRA + Qwen3-VL-4B text encoder
# + Qwen-Image VAE + Depth-Anything-V2-Large. Loaded once at module scope so
# ZeroGPU packs the weights and streams them to VRAM on the first GPU call.
print("Resolving Krea-2-Turbo base checkpoint...")
BASE_CKPT = os.path.realpath(hf_hub_download("krea/Krea-2-Turbo", "turbo.safetensors"))
LORA_CKPT = os.path.realpath(
    hf_hub_download("Patil/Krea-2-depth-controlnet", "depth-control-lora.safetensors")
)

print("Building DepthLoRAPipeline (13B DiT + Qwen3-VL-4B + VAE + DepthAnything)...")
pipe = DepthLoRAPipeline(BASE_CKPT, LORA_CKPT, device="cuda")
print("Pipeline ready.")


# --------------------------------------------------------------------- inference
@spaces.GPU(duration=120)
def generate(
    image,
    prompt: str = "",
    steps: int = 8,
    lora_scale: float = 1.0,
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a new image that keeps the 3D structure of an input image.

    Extracts a depth map from the input image with Depth-Anything-V2 and
    generates a new image following the same depth/composition but with the
    content and style described by the prompt (Krea-2-Turbo, 8-step).

    Args:
        image: The input image whose depth/structure is preserved.
        prompt: What to generate. Leave empty for depth-only generation.
        steps: Number of sampling steps (8 recommended for Turbo).
        lora_scale: Control strength. <1.0 relaxes structure adherence.
        seed: RNG seed for reproducibility.
        randomize_seed: If True, pick a random seed each run.

    Returns:
        A tuple of (generated image, extracted depth map, used seed).
    """
    if image is None:
        raise gr.Error("Please provide an input image.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    # Turbo config: cfg=0.0, mu=1.15. lora_scale is applied to the loaded LoRA
    # layers in place (they were built with scale=1.0), so scale the effective
    # weight by mutating each LoRALinear's scale before the run.
    from pipeline import LoRALinear

    for module in pipe.model.modules():
        if isinstance(module, LoRALinear):
            module.scale = (64 / 64) * float(lora_scale)

    out, depth = pipe(
        image,
        prompt=prompt or "",
        steps=int(steps),
        cfg=0.0,
        mu=1.15,
        seed=seed,
    )
    return out, depth, seed


# --------------------------------------------------------------------- UI
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # Krea-2 Depth ControlNet-LoRA

            Give it any image and a prompt — it extracts the depth map with
            **Depth-Anything-V2** and generates a new image with the **same 3D
            structure and composition**, but whatever content and style you ask
            for. Powered by [Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo)
            (8-step) + the
            [depth-control LoRA](https://huggingface.co/Patil/Krea-2-depth-controlnet).

            *Best with photos / renders that have real perspective. Flat 2D
            illustrations give weak control. Empty prompt = depth-only generation.*
            """
        )
        with gr.Row():
            with gr.Column():
                image = gr.Image(label="Input image (depth source)", type="pil")
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="a futuristic spaceship interior, cinematic lighting",
                    lines=2,
                )
                run = gr.Button("Generate", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    steps = gr.Slider(
                        4, 16, value=8, step=1, label="Sampling steps"
                    )
                    lora_scale = gr.Slider(
                        0.3,
                        1.4,
                        value=1.0,
                        step=0.05,
                        label="Control strength (LoRA scale)",
                    )
                    randomize_seed = gr.Checkbox(
                        label="Randomize seed", value=True
                    )
                    seed = gr.Slider(
                        0, MAX_SEED, value=0, step=1, label="Seed"
                    )
            with gr.Column():
                output = gr.Image(label="Generated image")
                depth_out = gr.Image(label="Extracted depth map")

        gr.Examples(
            examples=[
                ["examples/dog.jpg", "a majestic lion, golden hour, photorealistic"],
                [
                    "examples/landscape.jpg",
                    "an alien planet landscape, purple sky, sci-fi",
                ],
                [
                    "examples/man_beach.jpg",
                    "an astronaut on the moon, cinematic lighting",
                ],
                [
                    "examples/tent.jpg",
                    "a cozy cabin in a snowy forest at dusk",
                ],
            ],
            inputs=[image, prompt],
            outputs=[output, depth_out, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=generate,
        inputs=[image, prompt, steps, lora_scale, seed, randomize_seed],
        outputs=[output, depth_out, seed],
        api_name="generate",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)