import os
import subprocess
import sys

# ---------------------------------------------------------------------------
# LTX-2.4 prompt enhancer as a standalone ZeroGPU Space (Gemma-4 E2B). Kept
# separate so the 5GB enhancer + torchvision live off the main video Space's
# budget; the video Spaces call this over gradio_client. Replicates diffusers'
# LTX2Pipeline.enhance_prompt exactly (system prompt by mode, greedy decoding).
# The diffusers wheel is installed only for the LTX-2.4 system-prompt constants.
# ---------------------------------------------------------------------------
from huggingface_hub import hf_hub_download

HF_TOKEN = os.environ.get("HF_TOKEN")
_whl = hf_hub_download(
    "diffusers-internal-dev/ltx24-wheels",
    # rc2 wheel (PR head): identical LTX-2.4 system prompts, but GEMMA4_PROMPT_ENHANCEMENT_CONFIG
    # raises no_repeat_ngram_size 3 -> 5. At 3, long descriptive captions get pushed off common
    # phrasing and occasionally emit odd tokens (observed a Devanagari word spliced mid-caption).
    "rc2/diffusers-0.40.0.dev0-py3-none-any.whl",
    repo_type="dataset",
    token=HF_TOKEN,
)
_TARGET = "/tmp/ltx24_diffusers"
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--no-deps", "--target", _TARGET, _whl],
    check=True,
)
sys.path.insert(0, _TARGET)

import gradio as gr
import spaces
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

from diffusers.pipelines.ltx2.utils import (
    GEMMA4_PROMPT_ENHANCEMENT_CONFIG,
    LTX2_4_I2V_DEFAULT_SYSTEM_PROMPT,
    LTX2_4_T2V_DEFAULT_SYSTEM_PROMPT,
)

ENHANCER_ID = "google/gemma-4-E2B-it"
CFG = GEMMA4_PROMPT_ENHANCEMENT_CONFIG

print("[VERSION] ltx-2.4-enhancer v1 — loading", flush=True)
processor = AutoProcessor.from_pretrained(ENHANCER_ID, token=HF_TOKEN)
model = AutoModelForImageTextToText.from_pretrained(
    ENHANCER_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN
).to("cuda")
model.eval()
print("[VERSION] enhancer ready", flush=True)


@spaces.GPU(duration=120)
def enhance(prompt, image=None, progress=gr.Progress()):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt to enhance.")
    # I2V (reference-image) system prompt when an image is supplied, else T2V.
    system_prompt = LTX2_4_I2V_DEFAULT_SYSTEM_PROMPT if image is not None else LTX2_4_T2V_DEFAULT_SYSTEM_PROMPT
    user_text = f"{CFG.user_prompt_prefix}: {prompt}"
    # Gemma-4's processor validates that the rendered template carries exactly one <|image|> token per
    # image passed. A plain-string user turn renders none, so supplying an image raised
    #   ValueError: ... Found [0] <|image|> tokens and [1] images per sample
    # on EVERY image-to-video enhancement, while text-only worked — which is why it went unnoticed.
    # The two shapes below are exactly `LTX2Pipeline._enhance_prompt` (plain string, images=None) and
    # `LTX2ImageToVideoPipeline._enhance_prompt` (content blocks, images=image), so this Space keeps
    # returning what the pipelines' own enhancement returns. The text-only path is untouched.
    user_content = (
        [{"type": "image"}, {"type": "text", "text": user_text}] if image is not None else user_text
    )
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]
    template = processor.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    model_inputs = processor(text=template, images=image, return_tensors="pt").to("cuda")
    torch.manual_seed(10)  # greedy decoding is deterministic; seed is inert but matches diffusers
    with torch.no_grad():
        generated = model.generate(**model_inputs, max_new_tokens=512, **CFG.generation_kwargs)
    generated_ids = [seq[len(model_inputs.input_ids[i]):] for i, seq in enumerate(generated)]
    return processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]


with gr.Blocks(title="LTX-2.4 Prompt Enhancer") as demo:
    gr.Markdown(
        "# ✨ LTX-2.4 Prompt Enhancer\n"
        "Gemma-4 (`google/gemma-4-E2B-it`) prompt enhancer for LTX-2.4 — expands a short prompt into "
        "a detailed audio-visual caption in the model's training-caption style. Add a first-frame "
        "image to enhance for image-to-video. Called by the LTX-2.4 video Spaces over `gradio_client`."
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="e.g. a red fox in a snowy forest")
            image = gr.Image(label="First frame (optional — for image-to-video)", type="pil")
            btn = gr.Button("Enhance", variant="primary")
        out = gr.Textbox(label="Enhanced prompt", lines=12)
    btn.click(enhance, inputs=[prompt, image], outputs=[out], api_name="enhance")

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