| """Two-stage prompt rewriter for LingBot-Video (example-side helper). |
| |
| rewrite_prompt turns a brief idea into the structured JSON caption the DiT expects: |
| stage 1 expands the idea into a natural-language caption (base model), stage 2 maps it |
| into structured JSON (base model + stage-2 LoRA). This is a separate VLM + LoRA adapter, |
| not the DiT, and is not downloaded with the pipeline -- hence it lives with the examples. |
| TransformersBackend loads the rewriter VLM locally; make_backend also accepts any object |
| exposing generate(text, image, use_lora). |
| |
| from examples.lingbot_video.model_training.scripts.prompt_rewriter import rewrite_prompt |
| caption = rewrite_prompt("a puppy running across a meadow", mode="t2v", duration=5) |
| video = pipe(prompt=caption, ...) |
| """ |
|
|
| import io |
| import json |
| import os |
| import re |
|
|
| |
| try: |
| import requests |
| except ImportError: |
| requests = None |
|
|
| try: |
| from PIL import Image |
| except ImportError: |
| Image = None |
|
|
| try: |
| from json_repair import repair_json |
| except ImportError: |
| repair_json = None |
|
|
| try: |
| from .system_prompts import ( |
| VIDEO_STEP1_EXPAND, VIDEO_STEP2_MAP, IMAGE_STEP1_EXPAND, IMAGE_STEP2_MAP, |
| ) |
| except ImportError: |
| from system_prompts import ( |
| VIDEO_STEP1_EXPAND, VIDEO_STEP2_MAP, IMAGE_STEP1_EXPAND, IMAGE_STEP2_MAP, |
| ) |
|
|
|
|
| |
| MODES = { |
| "t2v": dict(s1=VIDEO_STEP1_EXPAND, s2=VIDEO_STEP2_MAP, image=False, duration=True), |
| "ti2v": dict(s1=VIDEO_STEP1_EXPAND, s2=VIDEO_STEP2_MAP, image=True, duration=True), |
| "t2i": dict(s1=IMAGE_STEP1_EXPAND, s2=IMAGE_STEP2_MAP, image=False, duration=False), |
| } |
|
|
|
|
| def _has_cjk(s: str) -> bool: |
| return any("一" <= c <= "鿿" for c in s) |
|
|
|
|
| def load_image(src): |
| """Load a first-frame image from a local path / http(s) URL / PIL.Image -> RGB PIL.""" |
| if Image is None: |
| raise ImportError("loading a first-frame image requires the Pillow package.") |
| if isinstance(src, Image.Image): |
| return src.convert("RGB") |
| if isinstance(src, str) and re.match(r"^https?://", src): |
| if requests is None: |
| raise ImportError("fetching an image URL requires the requests package.") |
| return Image.open(io.BytesIO(requests.get(src, timeout=30).content)).convert("RGB") |
| return Image.open(src).convert("RGB") |
|
|
|
|
| def _step1_text(mode, prompt, dur): |
| sys = MODES[mode]["s1"] |
| if mode == "t2i": |
| return sys + "\n\nUser image prompt:\n" + prompt |
| dur_line = f"\n\n视频时长：{dur} 秒" if _has_cjk(prompt) else f"\n\nVideo Duration: {dur} seconds" |
| return sys + "\n\n" + prompt + dur_line |
|
|
|
|
| def _step2_text(mode, detailed, dur): |
| sys = MODES[mode]["s2"] |
| if mode == "t2i": |
| return sys + "\n\nDETAILED CAPTION:\n" + detailed |
| return (sys + f"\n\nVideo Duration: {dur} seconds\n\nDETAILED CAPTION:\n" |
| + detailed + "\n\nOutput the JSON now.") |
|
|
|
|
| def parse_json(raw): |
| """Parse the stage-2 output into a dict, or None. Strips any code fence, tries the |
| stdlib parser, then falls back to json_repair (if installed) for messy output.""" |
| s = (raw or "").strip() |
| m = re.search(r"```(?:json)?\s*(\{.*\})\s*```", s, re.DOTALL) |
| if m: |
| s = m.group(1) |
| try: |
| obj = json.loads(s) |
| if isinstance(obj, dict): |
| return obj |
| except Exception: |
| pass |
| if repair_json is not None: |
| try: |
| obj = repair_json(s, return_objects=True) |
| return obj if isinstance(obj, dict) else None |
| except Exception: |
| return None |
| return None |
|
|
|
|
| def save_caption(result, duration, path): |
| """Save as {"caption": <caption>, "duration": <seconds|null>} -- the prompt.json the |
| pipeline consumes. duration is integer seconds for T2V/TI2V, None for T2I.""" |
| dur = int(round(duration)) if MODES[result["mode"]]["duration"] else None |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump({"caption": result["json"], "duration": dur}, f, ensure_ascii=False, indent=2) |
| return path |
|
|
|
|
| class TransformersBackend: |
| """Local rewriter VLM + LoRA adapter (peft): stage1 = base (adapter disabled), |
| stage2 = base + LoRA. Set base/adapter (or REWRITER_BASE_MODEL/REWRITER_ADAPTER).""" |
|
|
| def __init__(self, base=None, adapter=None, device="auto", max_new_tokens=6144): |
| import contextlib |
| import torch |
| from peft import PeftModel |
| from transformers import AutoModelForImageTextToText, AutoProcessor |
|
|
| self._contextlib = contextlib |
| self._torch = torch |
|
|
| base = base or os.environ.get("REWRITER_BASE_MODEL", "") |
| adapter = adapter or os.environ.get("REWRITER_ADAPTER", "") |
| if not base or not adapter: |
| raise ValueError( |
| "Set the rewriter base and adapter paths via base=/adapter= or the " |
| "REWRITER_BASE_MODEL / REWRITER_ADAPTER environment variables." |
| ) |
| self.processor = AutoProcessor.from_pretrained(base, trust_remote_code=True) |
| model = AutoModelForImageTextToText.from_pretrained( |
| base, torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True) |
| self.model = PeftModel.from_pretrained(model, adapter).eval() |
| self.max_new_tokens = max_new_tokens |
|
|
| def generate(self, text, image, use_lora): |
| torch = self._torch |
| content = ([{"type": "image", "image": image}] if image is not None else []) \ |
| + [{"type": "text", "text": text}] |
| messages = [{"role": "user", "content": content}] |
| chat = self.processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| inputs = self.processor( |
| text=[chat], images=([image] if image is not None else None), return_tensors="pt" |
| ).to(self.model.device) |
| |
| adapter_ctx = self._contextlib.nullcontext() if use_lora else self.model.disable_adapter() |
| with torch.no_grad(), adapter_ctx: |
| out = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens, do_sample=False) |
| gen = out[:, inputs["input_ids"].shape[1]:] |
| return self.processor.batch_decode(gen, skip_special_tokens=True)[0] |
|
|
|
|
| def make_backend(backend="transformers", base=None, adapter=None): |
| """Build a rewriter backend: "transformers" for the local TransformersBackend, or any |
| object already exposing generate(text, image, use_lora) -> str (e.g. a hosted endpoint).""" |
| if backend == "transformers": |
| return TransformersBackend(base, adapter) |
| if hasattr(backend, "generate"): |
| return backend |
| raise ValueError( |
| f"unknown backend: {backend!r}; pass 'transformers' or an object exposing " |
| "generate(text, image, use_lora)." |
| ) |
|
|
|
|
| class Rewriter: |
| """Two-stage orchestrator. The backend implements ``generate(text, image, use_lora) -> str``.""" |
|
|
| def __init__(self, backend): |
| self.backend = backend |
|
|
| def rewrite(self, prompt, mode="t2v", first_frame=None, duration=5): |
| if mode not in MODES: |
| raise ValueError(f"mode must be one of {list(MODES)}") |
| cfg = MODES[mode] |
| dur = int(round(duration)) |
| img = None |
| if cfg["image"]: |
| if first_frame is None: |
| raise ValueError(f"{mode} requires first_frame (path / URL / PIL.Image)") |
| img = load_image(first_frame) |
| |
| detailed = self.backend.generate(_step1_text(mode, prompt, dur), img, use_lora=False).strip() |
| |
| raw = self.backend.generate(_step2_text(mode, detailed, dur), img, use_lora=True).strip() |
| return {"mode": mode, "detailed": detailed, "json": parse_json(raw), "json_raw": raw} |
|
|
|
|
| def rewrite_prompt(prompt, mode="t2v", first_frame=None, duration=5, |
| backend="transformers", base=None, adapter=None, return_result=False): |
| """Rewrite a brief idea into the structured caption string the pipeline expects. |
| |
| Returns the compact-JSON caption string (ready to pass as ``pipe(prompt=...)``). |
| Pass ``return_result=True`` to also get the full stage-1/stage-2 dict. |
| """ |
| rw = Rewriter(make_backend(backend, base, adapter)) |
| result = rw.rewrite(prompt, mode=mode, first_frame=first_frame, duration=duration) |
| if result["json"] is not None: |
| caption = json.dumps(result["json"], ensure_ascii=False, separators=(",", ":")) |
| else: |
| caption = result["json_raw"] |
| if return_result: |
| return caption, result |
| return caption |
|
|
