Restore run.py with enhance_prompt binding (was overwritten by bad upload)

run.py

@@ -17,8 +17,11 @@ import os
 os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
 
+import hashlib
 import random
 import tempfile
+import time
+from collections import OrderedDict
 
 import gradio as gr
 import PIL.Image
@@ -247,6 +250,38 @@ def _generate_gpu(prompt, image_path, height, width, num_frames, seed, decoder,
     return path, num_frames
 
 
+# -----------------------------------------------------------------------------
+# Prompt enhancement — same mechanism and Space as the original demo (Gemma-4 on
+# a separate ZeroGPU Space). Authenticated with this Space's HF_TOKEN, so no
+# anonymous rate-limit bucket. Client rebuilt past a TTL: the ZeroGPU proxy
+# token it carries expires.
+# -----------------------------------------------------------------------------
+ENHANCER_SPACE = "diffusers-internal-dev/LTX-2.4-Prompt-Enhancer"
+_enh = {"client": None, "built_at": 0.0}
+_ENH_CLIENT_TTL_S = 900
+
+_ENH_CACHE: OrderedDict[tuple[str, str | None], str] = OrderedDict()
+_ENH_CACHE_MAX = 64
+
+
+def _image_key(path: str | None) -> str | None:
+    if not path:
+        return None
+    with open(path, "rb") as fh:
+        return hashlib.sha256(fh.read()).hexdigest()[:16]
+
+
+def _enhancer_client(fresh: bool = False):
+    from gradio_client import Client
+
+    if fresh or _enh["client"] is None or (time.monotonic() - _enh["built_at"]) > _ENH_CLIENT_TTL_S:
+        _enh["client"] = Client(
+            ENHANCER_SPACE, token=HF_TOKEN, httpx_kwargs={"timeout": 300.0}
+        )
+        _enh["built_at"] = time.monotonic()
+    return _enh["client"]
+
+
 def _extract_path(image_value):
     """Image ports arrive from the canvas as a {path|url} dict, a path string, or None."""
     if isinstance(image_value, dict):
@@ -254,6 +289,45 @@ def _extract_path(image_value):
     return image_value or None
 
 
+def enhance_prompt(prompt, image_path=None, do_enhance=True):
+    """CPU step (runs before any GPU is acquired): rewrite the prompt into a
+    detailed LTX-2.5-style caption via the enhancer Space. Falls back to the
+    raw prompt if the enhancer is unavailable. Bound as a `fn` operator."""
+    if not prompt or not str(prompt).strip():
+        raise gr.Error("Please enter a prompt.")
+    prompt = str(prompt).strip()
+    image_path = _extract_path(image_path)
+    if not do_enhance:
+        return prompt
+
+    key = (prompt, _image_key(image_path))
+    cached = _ENH_CACHE.get(key)
+    if cached is not None:
+        _ENH_CACHE.move_to_end(key)
+        return cached
+
+    from gradio_client import handle_file
+
+    args = (prompt, handle_file(image_path) if image_path else None)
+    try:
+        try:
+            enhanced = _enhancer_client().predict(*args, api_name="/enhance")
+        except Exception as first:  # noqa: BLE001
+            # Usually an expired ZeroGPU proxy token; retry once with a fresh client.
+            print(f"[enhance] first attempt failed ({first!r}); fresh client retry", flush=True)
+            enhanced = _enhancer_client(fresh=True).predict(*args, api_name="/enhance")
+    except Exception as e:  # noqa: BLE001
+        print(f"[enhance] failed, using raw prompt: {e!r}", flush=True)
+        gr.Warning(f"Prompt enhancement unavailable ({type(e).__name__}) — using your prompt as written.")
+        return prompt
+
+    enhanced = enhanced or prompt
+    _ENH_CACHE[key] = enhanced
+    while len(_ENH_CACHE) > _ENH_CACHE_MAX:
+        _ENH_CACHE.popitem(last=False)
+    return enhanced
+
+
 def generate_video(prompt, image_path=None, height=1472, width=832, duration_s=2,
                    seed=42, decoder="conv", auto_len=False, randomize_seed=True):
     """Workflow-facing wrapper, bound as a `fn` operator. CPU work (seed, frame
@@ -282,7 +356,7 @@ def generate_video(prompt, image_path=None, height=1472, width=832, duration_s=2
 
 demo = gr.Workflow(
     graph="workflow.json",
-    bind={"generate_video": generate_video},
+    bind={"enhance_prompt": enhance_prompt, "generate_video": generate_video},
 )
 
 if __name__ == "__main__":