"""Gradio Space: infinite-width panorama generation with Stable Diffusion v1.5
and the `infinite-tensor` framework.

This is a ZeroGPU wrapper around the paper's self-contained
``annotated_infinite_panorama.py`` demo script (Terrain Diffusion,
arXiv:2512.08309, https://xandergos.github.io/terrain-diffusion/). It
generates a seamless, arbitrarily-wide panorama by tiling latent diffusion
denoising across overlapping windows using the `infinite-tensor` library,
then VAE-decoding the result and cropping to the requested pixel width.

Only the infinite-panorama demo is wrapped here -- the hierarchical planetary
terrain pipeline and Flask API from the parent repo are intentionally NOT
ported (see BUILD_NOTES.md).
"""

import os

# --- 1. Cache env vars FIRST, before importing torch/diffusers/etc. ---------
# Prefer /data (persistent storage) if it's actually mounted and writable;
# otherwise fall back to a path under /home/user (always writable on
# ZeroGPU) so we never try to mkdir into a read-only/absent /data.
_data_cache = "/data/.cache/huggingface"
if os.path.isdir("/data") and os.access("/data", os.W_OK):
    os.environ.setdefault("HF_HOME", _data_cache)
else:
    os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import time

import spaces  # noqa: E402  (must import before torch)

import gradio as gr
import numpy as np
import torch
from diffusers import DDIMScheduler, StableDiffusionPipeline
from PIL import Image

from infinite_tensor import InfiniteTensor, MemoryTileStore, TensorWindow

# -----------------------------------------------------------------------------
# Model repo -- runwayml/stable-diffusion-v1-5 was removed from the Hub.
# Use the community mirror instead (same weights, still maintained).
# -----------------------------------------------------------------------------
MODEL_ID = "stable-diffusion-v1-5/stable-diffusion-v1-5"

# SDv1.5 constants (fixed by the pretrained UNet/VAE architecture).
LATENT_TILE = 64
PIXEL_TILE = 512
LATENT_CHANNELS = 4
LATENT_STRIDE = 32   # Overlap stride in latent space (tile = 64).
PIXEL_STRIDE = 384   # Overlap stride in pixel space (tile = 512).
INTERMEDIATE_TIMESTEPS = (400, 600, 750, 900)

# -----------------------------------------------------------------------------
# Load the pipeline once at import time. Module-level `.to("cuda")` is fine on
# ZeroGPU (CUDA is emulated at import time); only the actual forward passes
# need to happen inside a function decorated with @spaces.GPU.
# -----------------------------------------------------------------------------
pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    safety_checker=None,
).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)


# -----------------------------------------------------------------------------
# Helpers (verbatim logic from annotated_infinite_panorama.py)
# -----------------------------------------------------------------------------
def tiled_gaussian_noise(seed, x0, width, channels=LATENT_CHANNELS, height=LATENT_TILE, tile=256):
    """Sample a ``(channels, height, width)`` patch from a deterministic 1D-tiled
    Gaussian noise field. The value at column ``x`` depends only on
    ``(seed, x // tile)``, so overlapping tile requests agree, giving a
    self-consistent infinite noise field."""
    out = np.empty((channels, height, width), dtype=np.float32)
    first_tx = x0 // tile
    last_tx = (x0 + width - 1) // tile
    for tx in range(first_tx, last_tx + 1):
        tile_x0 = tx * tile
        ox0 = max(x0, tile_x0)
        ox1 = min(x0 + width, tile_x0 + tile)
        ss = np.random.SeedSequence(np.array([seed, tx & 0xFFFFFFFF], dtype=np.uint32))
        rng = np.random.Generator(np.random.PCG64DXSM(ss))
        tile_noise = rng.standard_normal((channels, height, tile), dtype=np.float32)
        out[:, :, ox0 - x0:ox1 - x0] = tile_noise[:, :, ox0 - tile_x0:ox1 - tile_x0]
    return out


def linear_kernel(height, width):
    """Separable linear blending weight, peak 1 at center, ~0 at edges."""
    x = torch.arange(width, dtype=torch.float32)
    mid = (width - 1) / 2
    w = 1 - 0.999 * torch.abs(x - mid) / mid
    return w[None, :].expand(height, -1).contiguous()


def build_timestep_ranges(all_timesteps, thresholds):
    """Partition descending ``all_timesteps`` into phases using ``thresholds``.
    Phase 0 gets ``t >= thresholds[0]``, the last phase gets
    ``t < thresholds[-1]``, intermediate phases fill the gaps."""
    thresholds = sorted(thresholds, reverse=True)
    if not thresholds:
        return [all_timesteps]
    ranges = []
    prev = None
    for t in thresholds:
        r = all_timesteps[all_timesteps >= t] if prev is None \
            else all_timesteps[(all_timesteps >= t) & (all_timesteps < prev)]
        if len(r) > 0:
            ranges.append(r)
        prev = t
    tail = all_timesteps[all_timesteps < thresholds[-1]]
    if len(tail) > 0:
        ranges.append(tail)
    return ranges


# -----------------------------------------------------------------------------
# Immutable blending kernels and tiling windows. These depend only on the fixed
# SDv1.5 tile geometry (constants above), so build them once at import time
# instead of reconstructing identical objects on every request.
# -----------------------------------------------------------------------------
LATENT_WEIGHT = linear_kernel(LATENT_TILE, LATENT_TILE)
PIXEL_WEIGHT = linear_kernel(PIXEL_TILE, PIXEL_TILE)
# DDIM init_noise_sigma is a scheduler constant (1.0); it does not depend on the
# per-request num_inference_steps, so it is safe to read once here.
INIT_NOISE_SIGMA = pipe.scheduler.init_noise_sigma

LATENT_WINDOW = TensorWindow(
    size=(LATENT_CHANNELS + 1, LATENT_TILE, LATENT_TILE),
    stride=(LATENT_CHANNELS + 1, LATENT_TILE, LATENT_STRIDE),
)
LATENT_DECODE_WINDOW = TensorWindow(
    size=(LATENT_CHANNELS + 1, LATENT_TILE, LATENT_TILE),
    stride=(LATENT_CHANNELS + 1, LATENT_TILE, PIXEL_STRIDE // 8),
)
PIXEL_WINDOW = TensorWindow(
    size=(3 + 1, PIXEL_TILE, PIXEL_TILE),
    stride=(3 + 1, PIXEL_TILE, PIXEL_STRIDE),
)


# -----------------------------------------------------------------------------
# Core generation, GPU-decorated.
# -----------------------------------------------------------------------------
@spaces.GPU(duration=120)
def generate_panorama(prompt, crop_width, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=False)):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a text prompt.")

    crop_width = int(crop_width)
    num_inference_steps = int(num_inference_steps)
    guidance_scale = float(guidance_scale)
    # Clamp into uint32 range: the tiled noise seeds a np.uint32 SeedSequence, so a
    # negative or oversized seed would overflow and crash the request.
    seed = int(seed) % (2**32)

    t_start = time.perf_counter()

    pipe.scheduler.set_timesteps(num_inference_steps)

    def encode(text):
        """Encode a string into CLIP text embeddings."""
        toks = pipe.tokenizer(
            text, padding="max_length", max_length=pipe.tokenizer.model_max_length,
            truncation=True, return_tensors="pt",
        ).input_ids.to(pipe.device)
        return pipe.text_encoder(toks)[0]

    # Text encoding needs no autograd graph (the UNet steps below already run
    # under no_grad); wrapping it here avoids tracking the CLIP encoder forward.
    with torch.no_grad():
        text_emb = torch.cat([encode(""), encode(prompt)])  # [uncond, cond]

    def denoise(latent, timesteps):
        """Run classifier-free-guided DDIM steps on a ``(1, C, H, W)`` latent."""
        for t in timesteps:
            inp = pipe.scheduler.scale_model_input(torch.cat([latent] * 2), t)
            with torch.no_grad():
                pred = pipe.unet(inp, t, encoder_hidden_states=text_emb).sample
            uncond, cond = pred.chunk(2)
            pred = uncond + guidance_scale * (cond - uncond)
            latent = pipe.scheduler.step(pred, t, latent).prev_sample
        return latent

    # phase_timesteps depends on num_inference_steps, so it stays per-request.
    phase_timesteps = build_timestep_ranges(pipe.scheduler.timesteps, INTERMEDIATE_TIMESTEPS)

    # Each latent/pixel tensor carries C+1 channels: C weighted values plus a
    # weight channel. infinite-tensor *sums* overlapping window outputs;
    # dividing the first C channels by the last recovers the weighted average
    # across overlapping tiles. That division is `normalize` below.
    def normalize(weighted):
        return weighted[:-1] / weighted[-1:].clamp(min=1e-6)

    def pack(values_chw, weight_hw):
        """``(C, H, W) + (H, W) -> (C+1, H, W)`` weighted output for infinite-tensor."""
        return torch.cat([values_chw * weight_hw[None], weight_hw[None]], dim=0)

    store = MemoryTileStore()
    T = len(phase_timesteps)

    def initial_phase(ctx):
        """Phase T-1: sample pure noise at this window column, denoise the highest-t range."""
        x = ctx[2] * LATENT_STRIDE
        noise = torch.as_tensor(tiled_gaussian_noise(seed, x, LATENT_TILE)) * INIT_NOISE_SIGMA
        noise = noise.to(pipe.device, dtype=torch.float16).unsqueeze(0)
        latent = denoise(noise, phase_timesteps[0])[0].cpu().float()
        return pack(latent, LATENT_WEIGHT)

    def make_continuation_phase(timesteps):
        """Phases T-2..0: read blended tile from previous phase, denoise further."""
        def phase(ctx, prev):
            latent = normalize(prev).to(pipe.device, dtype=torch.float16).unsqueeze(0)
            latent = denoise(latent, timesteps)[0].cpu().float()
            return pack(latent, LATENT_WEIGHT)
        return phase

    def decode(ctx, prev):
        """VAE-decode the fully denoised (phase 0) latent tile; re-weight for pixel blending."""
        latent = normalize(prev).to(pipe.device, dtype=torch.float16).unsqueeze(0)
        latent = latent / pipe.vae.config.scaling_factor
        with torch.no_grad():
            img = pipe.vae.decode(latent).sample
        img = (img / 2 + 0.5).clamp(0, 1)[0].cpu().float()
        return pack(img, PIXEL_WEIGHT)

    latents = InfiniteTensor(
        shape=(LATENT_CHANNELS + 1, LATENT_TILE, None),
        f=initial_phase,
        output_window=LATENT_WINDOW,
        tile_store=store,
        tensor_id=f"phase{T - 1}",
    )
    for i, timesteps in enumerate(phase_timesteps[1:], start=1):
        latents = InfiniteTensor(
            shape=(LATENT_CHANNELS + 1, LATENT_TILE, None),
            f=make_continuation_phase(timesteps),
            output_window=LATENT_WINDOW,
            args=(latents,),
            args_windows=(LATENT_WINDOW,),
            tile_store=store,
            tensor_id=f"phase{T - 1 - i}",
        )
    pixels = InfiniteTensor(
        shape=(3 + 1, PIXEL_TILE, None),
        f=decode,
        output_window=PIXEL_WINDOW,
        args=(latents,),
        args_windows=(LATENT_DECODE_WINDOW,),
        tile_store=store,
        tensor_id="image",
    )

    region = normalize(torch.as_tensor(pixels[:, :, 0:crop_width]))
    arr = (region.permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8)
    image = Image.fromarray(arr)

    elapsed = time.perf_counter() - t_start
    print(f"[terrain-diffusion] generated {crop_width}px panorama in {elapsed:.1f}s "
          f"(steps={num_inference_steps}, guidance={guidance_scale}, seed={seed})")

    return image


# -----------------------------------------------------------------------------
# UI
# -----------------------------------------------------------------------------
TITLE = "Terrain Diffusion — Infinite Panorama"

DESCRIPTION = """
Generate a seamless, arbitrarily-wide panorama with Stable Diffusion v1.5 and
the [`infinite-tensor`](https://github.com/xandergos/infinite-tensor) library.

This demo wraps the paper's self-contained `annotated_infinite_panorama.py`
script: noise is sampled from a deterministic tiled Gaussian field, denoised
in overlapping windows across several timestep phases (blended with a linear
kernel), then VAE-decoded and cropped to the requested width. No
outpainting/stitching artifacts — the panorama is coherent because every
overlapping tile is denoised from the *same* underlying infinite noise field.

**Paper:** Terrain Diffusion ([arXiv:2512.08309](https://arxiv.org/abs/2512.08309)) &middot;
[project page](https://xandergos.github.io/terrain-diffusion/)

Note: this Space demonstrates only the flat infinite-panorama demo from the
paper's repository, not the full hierarchical planetary-terrain pipeline.
"""

ARTICLE = """
### Citation
```bibtex
@inproceedings{goslin2026infinitediffusion,
      author    = {Goslin, Alexander},
      title     = {InfiniteDiffusion: Bridging Learned Fidelity and Procedural Utility for Open-World Terrain Generation},
      booktitle = {Special Interest Group on Computer Graphics and Interactive Techniques Conference Conference Papers},
      year      = {2026},
      pages     = {10 pages},
      publisher = {ACM},
      address   = {New York, NY, USA},
      doi       = {10.1145/3799902.3811080},
      url       = {https://doi.org/10.1145/3799902.3811080},
      series    = {SIGGRAPH Conference Papers '26}
}
```
"""

with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                value="a photo of a mountain range at sunset",
                placeholder="e.g. a photo of a mountain range at sunset",
                lines=2,
            )
            crop_width = gr.Slider(
                label="Crop width (px)",
                minimum=512, maximum=2048, step=128, value=1536,
                info="Width of the final panorama. Larger = more overlapping tiles = slower.",
            )
            steps = gr.Slider(
                label="Inference steps", minimum=10, maximum=100, step=1, value=40,
            )
            guidance = gr.Slider(
                label="Guidance scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5,
            )
            seed = gr.Number(label="Seed", value=0, precision=0)
            run_btn = gr.Button("Generate panorama", variant="primary")

        with gr.Column(scale=2):
            output_image = gr.Image(label="Panorama", type="pil")

    gr.Examples(
        examples=[
            ["a photo of a mountain range at sunset", 1536, 40, 7.5, 0],
            ["an oil painting of rolling green hills under a stormy sky", 1536, 40, 7.5, 1],
            ["a satellite photo of a desert canyon landscape", 1536, 40, 7.5, 2],
        ],
        inputs=[prompt, crop_width, steps, guidance, seed],
        outputs=output_image,
        fn=generate_panorama,
        cache_examples=False,
    )

    gr.Markdown(ARTICLE)

    run_btn.click(
        fn=generate_panorama,
        inputs=[prompt, crop_width, steps, guidance, seed],
        outputs=output_image,
        api_name="generate",
    )

if __name__ == "__main__":
    demo.queue().launch()