179 lines
7.8 KiB
Python
179 lines
7.8 KiB
Python
r"""Cache Qwen3-VL conditioning for data-free PDD (`fl2va`), multi-GPU, in safetensors.
|
|
|
|
Rewritten from the single-GPU `encode_prompts.py` to follow the repository's preprocessing convention
|
|
(mirrored on `scripts/wan2.1_self_forcing/generate_ode_pairs.py`):
|
|
`accelerate launch` over every rank, interleaved rank sharding, resume on files that already exist, per-entry
|
|
`.safetensors` (never `.pt` / LMDB), `wait_for_everyone`, then rank0 writes an `outputs.json` that
|
|
`ImageVideoSafetensorsDataset` consumes. Each entry holds the exact keys `FL2VATrajectory.reset` reads —
|
|
`prompt_embeds` (`hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER]`) and `text_token_tags` — so the ~62 GB conditioner
|
|
stays out of the PDD training run under `--enable_preprocess_training`.
|
|
|
|
accelerate launch --mixed_precision="bf16" scripts/minimax_h3/generate_prompt_cache.py \
|
|
--pretrained_model_name_or_path=models/Diffusion_Transformer/MiniMax-H3 \
|
|
--train_data_meta=datasets/X-Fun-Videos-Audios-Demo/metadata_add_width_height.json \
|
|
--output_folder=datasets/minimax_h3_pdd_prompt_cache
|
|
|
|
`--train_data_meta` is the official demo dataset's annotation JSON (never an ad-hoc prompt set). PDD is
|
|
data-free, so only its `text` captions are read — the `file_path` / `audio_path` / `control_file_path` / `width` /
|
|
`height` fields the audio-visual dataset carries are ignored here. `load_prompts` also accepts a bare list of prompt
|
|
strings or `{"prompt": ...}` records. Data-free PDD has no train/val split: cache once and point both
|
|
`--train_data_meta` and `--val_data_meta` of `train_pdd_lora.py` at the same `outputs.json`.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
from accelerate import Accelerator
|
|
from safetensors.torch import save_file
|
|
from tqdm import tqdm
|
|
|
|
current_file_path = os.path.abspath(__file__)
|
|
project_roots = [
|
|
os.path.dirname(current_file_path),
|
|
os.path.dirname(os.path.dirname(current_file_path)),
|
|
os.path.dirname(os.path.dirname(os.path.dirname(current_file_path))),
|
|
]
|
|
for project_root in project_roots:
|
|
sys.path.insert(0, project_root) if project_root not in sys.path else None
|
|
|
|
from videox_fun.models import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration
|
|
|
|
# Reuse the canonical MiniMax-H3 conditioner recipe rather than re-deriving it: `train_lora.encode_prompt` builds the
|
|
# presentation and reads `hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER]` exactly as the pipeline does. For a text-only
|
|
# `fl2va` request it never touches `processor`, so `None` is passed.
|
|
from train_lora import encode_prompt
|
|
|
|
|
|
def load_prompts(path):
|
|
r"""Read the annotation JSON into a list of prompt strings.
|
|
|
|
Accepts the repository annotation format (a list of `{"text": ...}`), a bare list of strings, and the legacy
|
|
`{"prompt": ...}` job records / `{"examples": [...]}` wrapper the old `encode_prompts.py` took.
|
|
"""
|
|
with open(path, encoding="utf-8") as handle:
|
|
document = json.load(handle)
|
|
if isinstance(document, dict):
|
|
document = document.get("examples", document)
|
|
if not isinstance(document, list) or not document:
|
|
raise ValueError(f"{path} must be a non-empty list of prompts or `{{'text': ...}}` records.")
|
|
prompts = []
|
|
for index, entry in enumerate(document, start=1):
|
|
if isinstance(entry, str):
|
|
prompt = entry.strip()
|
|
elif isinstance(entry, dict) and isinstance(entry.get("text"), str):
|
|
prompt = entry["text"].strip()
|
|
elif isinstance(entry, dict) and isinstance(entry.get("prompt"), str):
|
|
prompt = entry["prompt"].strip()
|
|
else:
|
|
raise ValueError(f"Entry {index} in {path} is not a prompt string or a `{{'text'/'prompt': ...}}` record.")
|
|
if not prompt:
|
|
raise ValueError(f"Entry {index} in {path} is empty.")
|
|
prompts.append(prompt)
|
|
return prompts
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Cache Qwen3-VL `fl2va` conditioning for data-free PDD (multi-GPU).")
|
|
parser.add_argument(
|
|
"--pretrained_model_name_or_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to the MiniMax-H3 partition; its `tokenizer/` and `text_encoder/` subfolders are read.",
|
|
)
|
|
parser.add_argument(
|
|
"--train_data_meta",
|
|
type=str,
|
|
required=True,
|
|
help="Annotation JSON of prompts: a list of `{\"text\": ...}` records (bare strings / `{\"prompt\": ...}` also accepted).",
|
|
)
|
|
parser.add_argument(
|
|
"--output_folder",
|
|
type=str,
|
|
required=True,
|
|
help="Directory the per-prompt `.safetensors` and the rank0 `outputs.json` are written to (one split).",
|
|
)
|
|
parser.add_argument(
|
|
"--mixed_precision",
|
|
type=str,
|
|
default="bf16",
|
|
choices=["no", "fp16", "bf16"],
|
|
help="Mixed precision the conditioner runs at. MiniMax-H3 conditions in bfloat16. Default: bf16.",
|
|
)
|
|
parser.add_argument("--local_rank", type=int, default=-1, help="For distributed preprocessing: local_rank.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
accelerator = Accelerator(mixed_precision=args.mixed_precision)
|
|
device = accelerator.device
|
|
world_size = accelerator.num_processes
|
|
rank = accelerator.process_index
|
|
|
|
torch.set_grad_enabled(False)
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
|
|
weight_dtype = torch.float32
|
|
if accelerator.mixed_precision == "fp16":
|
|
weight_dtype = torch.float16
|
|
elif accelerator.mixed_precision == "bf16":
|
|
weight_dtype = torch.bfloat16
|
|
|
|
prompts = load_prompts(args.train_data_meta)
|
|
tokenizer = Qwen2TokenizerFast.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "tokenizer"))
|
|
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, "text_encoder"), low_cpu_mem_usage=True, torch_dtype=weight_dtype,
|
|
).to(device).eval()
|
|
text_encoder.requires_grad_(False)
|
|
|
|
os.makedirs(args.output_folder, exist_ok=True)
|
|
total_per_rank = int(math.ceil(len(prompts) / world_size))
|
|
|
|
# Each rank walks an interleaved slice of the prompt list; a file that already exists is skipped so a killed run
|
|
# resumes instead of recomputing.
|
|
for index in tqdm(range(total_per_rank), disable=rank != 0, desc="Caching fl2va prompts"):
|
|
prompt_index = index * world_size + rank
|
|
if prompt_index >= len(prompts):
|
|
continue
|
|
prompt = prompts[prompt_index]
|
|
output_path = os.path.join(args.output_folder, f"{prompt_index:05d}.safetensors")
|
|
if os.path.exists(output_path):
|
|
continue
|
|
|
|
prompt_embeds, text_token_tags = encode_prompt(
|
|
text_encoder, tokenizer, None, prompt, device=device, dtype=weight_dtype,
|
|
)
|
|
save_file(
|
|
{
|
|
"prompt_embeds": prompt_embeds.to(torch.bfloat16).cpu().contiguous(),
|
|
"text_token_tags": text_token_tags.cpu().contiguous().long(),
|
|
},
|
|
output_path,
|
|
metadata={"format": "pt", "prompt": prompt},
|
|
)
|
|
|
|
accelerator.wait_for_everyone()
|
|
|
|
# rank0 lists every generated file so `ImageVideoSafetensorsDataset` can load the split.
|
|
if accelerator.is_main_process:
|
|
records = []
|
|
for prompt_index in range(len(prompts)):
|
|
safetensor_path = os.path.join(args.output_folder, f"{prompt_index:05d}.safetensors")
|
|
if os.path.exists(safetensor_path):
|
|
records.append({"file_path": safetensor_path})
|
|
json_path = os.path.join(args.output_folder, "outputs.json")
|
|
with open(json_path, "w", encoding="utf-8") as handle:
|
|
json.dump(records, handle, ensure_ascii=False, indent=4)
|
|
print(f"Done. Cached {len(records)} fl2va prompts to {args.output_folder}")
|
|
print(f"Annotation JSON: {json_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|