Spaces:
Running on Zero
Running on Zero
File size: 3,804 Bytes
e487c51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | 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",
"canon/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
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"{CFG.user_prompt_prefix}: {prompt}"},
]
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, show_copy_button=True)
btn.click(enhance, inputs=[prompt, image], outputs=[out], api_name="enhance")
if __name__ == "__main__":
demo.launch(show_error=True)
|