Download scripts/generate_vbench8_extended_naive_baselines.py from Cccccz/Self-Forcing: direct link, hf CLI and curl.
- Browser
- Download file 22.7 kB
-
https://huggingface.co/Cccccz/Self-Forcing/resolve/main/scripts/generate_vbench8_extended_naive_baselines.py
- Command line
-
hf download hf://Cccccz/Self-Forcing/scripts/generate_vbench8_extended_naive_baselines.py
-
curl -L -o generate_vbench8_extended_naive_baselines.py https://huggingface.co/Cccccz/Self-Forcing/resolve/main/scripts/generate_vbench8_extended_naive_baselines.py
22.7 kB
| #!/usr/bin/env python3 | |
| """Generate Extended-251 videos for timestep-skipping and velocity-cache baselines. | |
| For every prompt, an original four-step FFFF rollout is generated first with | |
| the same seed and used as the PSNR/SSIM/LPIPS reference. Timestep-skipping | |
| policies shorten the denoising schedule itself. Velocity-cache policies retain | |
| the original four timesteps and, at each R, reuse the velocity from the most | |
| recent F while converting it to x0 using the current noisy latent and timestep. | |
| The cache pattern is applied independently to every autoregressive chunk. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import sys | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| from typing import Any | |
| def preparse_gpu() -> str: | |
| parser = argparse.ArgumentParser(add_help=False) | |
| parser.add_argument("--gpu", default="0") | |
| args, _ = parser.parse_known_args() | |
| os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) | |
| os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="self_forcing_mpl_")) | |
| return str(args.gpu) | |
| PHYSICAL_GPU = preparse_gpu() | |
| import lpips | |
| import torch | |
| from omegaconf import OmegaConf | |
| REPO_ROOT = Path(__file__).resolve().parents[1] | |
| if str(REPO_ROOT) not in sys.path: | |
| sys.path.insert(0, str(REPO_ROOT)) | |
| from scripts import evaluate_single_block_fppf as base | |
| from scripts.naive_vbench_policies import ( | |
| EVALUATION_STRATEGY_NAMES, | |
| POLICIES_BY_NAME, | |
| REFERENCE_POLICY, | |
| NaivePolicy, | |
| ) | |
| from scripts.vbench8_protocol import SUITE_COUNTS | |
| from utils.misc import set_seed | |
| from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper | |
| NUM_CHUNKS = base.NUM_CHUNKS | |
| FRAMES_PER_CHUNK = base.FRAMES_PER_CHUNK | |
| MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json" | |
| OUTPUT_DEFAULT = REPO_ROOT / "evaluation_runs/vbench8_extended_naive_baselines" | |
| def atomic_json(path: Path, value: Any) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = path.with_suffix(path.suffix + ".tmp") | |
| temporary.write_text( | |
| json.dumps(value, ensure_ascii=False, indent=2, allow_nan=True) + "\n", | |
| encoding="utf-8", | |
| ) | |
| os.replace(temporary, path) | |
| def atomic_video(frames: torch.Tensor, path: Path) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = path.with_name(path.name + ".partial.mp4") | |
| if temporary.exists(): | |
| temporary.unlink() | |
| base.save_mp4(frames, temporary) | |
| os.replace(temporary, path) | |
| def read_mapping(path: Path) -> list[dict[str, Any]]: | |
| value = json.loads(path.read_text(encoding="utf-8")) | |
| if not isinstance(value, list) or len(value) != 251: | |
| raise ValueError(f"Expected a 251-row mapping list: {path}") | |
| counts = {suite: 0 for suite in SUITE_COUNTS} | |
| globals_seen: set[int] = set() | |
| suite_seen: dict[str, set[int]] = {suite: set() for suite in SUITE_COUNTS} | |
| for row in value: | |
| suite = str(row["prompt_suite"]) | |
| global_index = int(row["global_index"]) | |
| suite_index = int(row["suite_index"]) | |
| if suite not in counts: | |
| raise ValueError(f"Unknown prompt suite: {suite}") | |
| if global_index in globals_seen or suite_index in suite_seen[suite]: | |
| raise ValueError( | |
| f"Duplicate mapping index: {global_index} / {suite}:{suite_index}" | |
| ) | |
| if not str(row["extended_prompt"]).strip(): | |
| raise ValueError(f"Empty prompt at global index {global_index}") | |
| counts[suite] += 1 | |
| globals_seen.add(global_index) | |
| suite_seen[suite].add(suite_index) | |
| if counts != SUITE_COUNTS: | |
| raise ValueError(f"Unexpected prompt-suite counts: {counts}") | |
| return sorted(value, key=lambda row: int(row["global_index"])) | |
| def reset_runtime_caches(pipeline: Any, device: torch.device) -> None: | |
| if pipeline.kv_cache1 is None: | |
| pipeline._initialize_kv_cache(1, torch.bfloat16, device) | |
| pipeline._initialize_crossattn_cache(1, torch.bfloat16, device) | |
| for cache in pipeline.kv_cache1: | |
| cache["global_end_index"].zero_() | |
| cache["local_end_index"].zero_() | |
| for cache in pipeline.crossattn_cache: | |
| cache["is_init"] = False | |
| def warped_timesteps(pipeline: Any, raw_timesteps: tuple[int, ...]) -> torch.Tensor: | |
| raw = torch.tensor(raw_timesteps, dtype=torch.long) | |
| if not bool(pipeline.args.warp_denoising_step): | |
| return raw | |
| scheduler_timesteps = torch.cat( | |
| (pipeline.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)) | |
| ) | |
| return scheduler_timesteps[1000 - raw] | |
| def policy_timesteps(pipeline: Any, policy: NaivePolicy) -> torch.Tensor: | |
| if policy.family in {"full", "velocity_reuse"}: | |
| actual = pipeline.denoising_step_list.detach().cpu() | |
| else: | |
| actual = warped_timesteps(pipeline, policy.raw_timesteps) | |
| if len(actual) != policy.denoise_steps: | |
| raise ValueError( | |
| f"{policy.name}: expected {policy.denoise_steps} timesteps, got {len(actual)}" | |
| ) | |
| return actual | |
| def start_timing() -> tuple[torch.cuda.Event, torch.cuda.Event]: | |
| start = torch.cuda.Event(enable_timing=True) | |
| end = torch.cuda.Event(enable_timing=True) | |
| start.record() | |
| return start, end | |
| def finish_timing( | |
| events_by_category: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]], | |
| category: str, | |
| events: tuple[torch.cuda.Event, torch.cuda.Event], | |
| ) -> None: | |
| events[1].record() | |
| events_by_category[category].append(events) | |
| def generate_rollout( | |
| *, | |
| pipeline: Any, | |
| conditional_dict: dict[str, torch.Tensor], | |
| seed: int, | |
| device: torch.device, | |
| policy: NaivePolicy, | |
| ) -> tuple[torch.Tensor, dict[str, Any]]: | |
| reset_runtime_caches(pipeline, device) | |
| set_seed(seed) | |
| noise = torch.randn( | |
| 1, | |
| NUM_CHUNKS * FRAMES_PER_CHUNK, | |
| base.LATENT_CHANNELS, | |
| base.LATENT_HEIGHT, | |
| base.LATENT_WIDTH, | |
| dtype=torch.bfloat16, | |
| device=device, | |
| ) | |
| timesteps_cpu = policy_timesteps(pipeline, policy) | |
| actual_timestep_values = [float(value) for value in timesteps_cpu] | |
| timesteps = timesteps_cpu.to(device=device) | |
| output_chunks: list[torch.Tensor] = [] | |
| decisions: list[dict[str, Any]] = [] | |
| full_calls = 0 | |
| reuse_calls = 0 | |
| timing_events: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = { | |
| "full_dit": [], | |
| "context_dit": [], | |
| } | |
| started = time.perf_counter() | |
| for chunk in range(NUM_CHUNKS): | |
| noisy_input = noise[ | |
| :, chunk * FRAMES_PER_CHUNK : (chunk + 1) * FRAMES_PER_CHUNK | |
| ] | |
| cached_flow: torch.Tensor | None = None | |
| denoised_pred: torch.Tensor | None = None | |
| timestep: torch.Tensor | None = None | |
| for step, current_timestep in enumerate(timesteps): | |
| timestep = torch.ones( | |
| [1, FRAMES_PER_CHUNK], dtype=torch.int64, device=device | |
| ) * current_timestep | |
| action = ( | |
| policy.cache_pattern[step] | |
| if policy.family == "velocity_reuse" and policy.cache_pattern is not None | |
| else "F" | |
| ) | |
| if action == "R": | |
| if cached_flow is None: | |
| raise RuntimeError( | |
| f"{policy.name}: reuse requested before a full step in chunk {chunk}" | |
| ) | |
| denoised_pred = pipeline.generator._convert_flow_pred_to_x0( | |
| flow_pred=cached_flow.flatten(0, 1), | |
| xt=noisy_input.flatten(0, 1), | |
| timestep=timestep.flatten(0, 1), | |
| ).unflatten(0, cached_flow.shape[:2]) | |
| reuse_calls += 1 | |
| else: | |
| full_events = start_timing() | |
| flow, denoised_pred = pipeline.generator( | |
| noisy_image_or_video=noisy_input, | |
| conditional_dict=conditional_dict, | |
| timestep=timestep, | |
| kv_cache=pipeline.kv_cache1, | |
| crossattn_cache=pipeline.crossattn_cache, | |
| current_start=chunk * base.TOKENS_PER_CHUNK, | |
| ) | |
| finish_timing(timing_events, "full_dit", full_events) | |
| cached_flow = flow.detach() | |
| full_calls += 1 | |
| decisions.append( | |
| { | |
| "chunk": chunk, | |
| "step": step, | |
| "action": action, | |
| "timestep": actual_timestep_values[step], | |
| } | |
| ) | |
| if step < len(timesteps) - 1: | |
| if denoised_pred is None: | |
| raise RuntimeError("Denoising step did not produce x0") | |
| next_timestep = timesteps[step + 1] | |
| flat = denoised_pred.flatten(0, 1) | |
| noisy_input = pipeline.scheduler.add_noise( | |
| flat, | |
| torch.randn_like(flat), | |
| next_timestep | |
| * torch.ones( | |
| [FRAMES_PER_CHUNK], dtype=torch.long, device=device | |
| ), | |
| ).unflatten(0, denoised_pred.shape[:2]) | |
| if denoised_pred is None or timestep is None: | |
| raise RuntimeError(f"{policy.name}: chunk {chunk} produced no clean latent") | |
| output_chunks.append(denoised_pred) | |
| context_events = start_timing() | |
| pipeline.generator( | |
| noisy_image_or_video=denoised_pred, | |
| conditional_dict=conditional_dict, | |
| timestep=torch.ones_like(timestep) * pipeline.args.context_noise, | |
| kv_cache=pipeline.kv_cache1, | |
| crossattn_cache=pipeline.crossattn_cache, | |
| current_start=chunk * base.TOKENS_PER_CHUNK, | |
| ) | |
| finish_timing(timing_events, "context_dit", context_events) | |
| torch.cuda.synchronize() | |
| elapsed = { | |
| category: sum(start.elapsed_time(end) for start, end in events) | |
| for category, events in timing_events.items() | |
| } | |
| expected_full = NUM_CHUNKS * policy.full_calls_per_chunk | |
| expected_reuse = NUM_CHUNKS * policy.reuse_calls_per_chunk | |
| if (full_calls, reuse_calls) != (expected_full, expected_reuse): | |
| raise RuntimeError( | |
| f"{policy.name}: got calls F={full_calls}, R={reuse_calls}; " | |
| f"expected F={expected_full}, R={expected_reuse}" | |
| ) | |
| return torch.cat(output_chunks, dim=1), { | |
| "generation_time_s": time.perf_counter() - started, | |
| "full_calls": full_calls, | |
| "reuse_calls": reuse_calls, | |
| "predictor_calls": 0, | |
| "accepted_predictor_calls": 0, | |
| "rejected_predictor_calls": 0, | |
| "full_dit_time_ms": elapsed["full_dit"], | |
| "predictor_time_ms": 0.0, | |
| "confidence_head_time_ms": 0.0, | |
| "context_dit_time_ms": elapsed["context_dit"], | |
| "actual_dit_time_ms": elapsed["full_dit"] + elapsed["context_dit"], | |
| "model_path_time_ms": elapsed["full_dit"], | |
| "raw_timesteps": list(policy.raw_timesteps), | |
| "actual_timesteps": actual_timestep_values, | |
| "decisions": decisions, | |
| } | |
| def compact_metrics(value: dict[str, Any]) -> dict[str, Any]: | |
| keys = ( | |
| "pixel_mse", | |
| "psnr", | |
| "psnr_frame_mean", | |
| "ssim", | |
| "lpips", | |
| "rollout_start_frame", | |
| "rollout_pixel_mse", | |
| "rollout_psnr", | |
| "rollout_ssim", | |
| "rollout_lpips", | |
| "num_frames", | |
| ) | |
| return {key: value[key] for key in keys} | |
| def identity_metrics(num_frames: int) -> dict[str, Any]: | |
| return { | |
| "pixel_mse": 0.0, | |
| "psnr": 120.0, | |
| "psnr_frame_mean": 120.0, | |
| "ssim": 1.0, | |
| "lpips": 0.0, | |
| "rollout_start_frame": base.PIXEL_FRAMES_FIRST_CHUNK, | |
| "rollout_pixel_mse": 0.0, | |
| "rollout_psnr": 120.0, | |
| "rollout_ssim": 1.0, | |
| "rollout_lpips": 0.0, | |
| "num_frames": num_frames, | |
| } | |
| def load_models(device: torch.device) -> tuple[Any, Any, Any, Any]: | |
| print("[setup] loading VAE", flush=True) | |
| vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval() | |
| config = OmegaConf.merge( | |
| OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"), | |
| OmegaConf.load(REPO_ROOT / "configs/self_forcing_sid.yaml"), | |
| ) | |
| print("[setup] loading frozen Self-Forcing Teacher", flush=True) | |
| pipeline = base.build_pipeline( | |
| config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device | |
| ) | |
| text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval() | |
| text_encoder.requires_grad_(False) | |
| print("[setup] loading LPIPS", flush=True) | |
| lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval() | |
| lpips_model.requires_grad_(False) | |
| return vae, pipeline, text_encoder, lpips_model | |
| def artifact_paths( | |
| output_root: Path, mapping_row: dict[str, Any], strategy: str | |
| ) -> tuple[Path, Path]: | |
| suite = str(mapping_row["prompt_suite"]) | |
| suite_index = int(mapping_row["suite_index"]) | |
| global_index = int(mapping_row["global_index"]) | |
| return ( | |
| output_root / "generated_videos" / strategy / suite / f"{suite_index:03d}.mp4", | |
| output_root | |
| / "generation_metrics/per_prompt" | |
| / strategy | |
| / f"global_{global_index:04d}.json", | |
| ) | |
| def is_prompt_complete( | |
| output_root: Path, mapping_row: dict[str, Any], policies: tuple[NaivePolicy, ...] | |
| ) -> bool: | |
| for policy in (REFERENCE_POLICY, *policies): | |
| video, record = artifact_paths(output_root, mapping_row, policy.name) | |
| if not video.is_file() or not record.is_file(): | |
| return False | |
| return True | |
| def write_result( | |
| *, | |
| output_root: Path, | |
| mapping_row: dict[str, Any], | |
| policy: NaivePolicy, | |
| seed: int, | |
| diagnostic: dict[str, Any], | |
| pixel_metrics: dict[str, Any], | |
| frames: torch.Tensor, | |
| ) -> None: | |
| video_path, record_path = artifact_paths(output_root, mapping_row, policy.name) | |
| atomic_video(frames, video_path) | |
| diagnostic_without_decisions = { | |
| key: value for key, value in diagnostic.items() if key != "decisions" | |
| } | |
| record = { | |
| "status": "complete", | |
| "strategy": policy.name, | |
| "label": policy.label, | |
| "policy": policy.to_dict(), | |
| "global_index": int(mapping_row["global_index"]), | |
| "prompt_suite": str(mapping_row["prompt_suite"]), | |
| "suite_index": int(mapping_row["suite_index"]), | |
| "prompt": str(mapping_row["extended_prompt"]), | |
| "seed": seed, | |
| "generation": diagnostic_without_decisions, | |
| "decisions": diagnostic["decisions"], | |
| "pixel_metrics_vs_ffff": pixel_metrics, | |
| "video": str(video_path.relative_to(output_root)), | |
| } | |
| atomic_json(record_path, record) | |
| def run_prompt( | |
| *, | |
| mapping_row: dict[str, Any], | |
| policies: tuple[NaivePolicy, ...], | |
| output_root: Path, | |
| seed: int, | |
| vae: Any, | |
| pipeline: Any, | |
| text_encoder: Any, | |
| lpips_model: Any, | |
| device: torch.device, | |
| ) -> None: | |
| global_index = int(mapping_row["global_index"]) | |
| suite = str(mapping_row["prompt_suite"]) | |
| suite_index = int(mapping_row["suite_index"]) | |
| prompt = str(mapping_row["extended_prompt"]) | |
| print(f"[prompt] global={global_index} suite={suite}/{suite_index}", flush=True) | |
| conditional = text_encoder(text_prompts=[prompt]) | |
| reference_latent, reference_diagnostic = generate_rollout( | |
| pipeline=pipeline, | |
| conditional_dict=conditional, | |
| seed=seed, | |
| device=device, | |
| policy=REFERENCE_POLICY, | |
| ) | |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16): | |
| reference_pixels = vae.decode_to_pixel(reference_latent, use_cache=False) | |
| reference_u8 = base.pixels_to_u8(reference_pixels) | |
| if reference_u8.ndim != 4 or reference_u8.shape[0] != 81: | |
| raise RuntimeError(f"Unexpected reference shape: {tuple(reference_u8.shape)}") | |
| write_result( | |
| output_root=output_root, | |
| mapping_row=mapping_row, | |
| policy=REFERENCE_POLICY, | |
| seed=seed, | |
| diagnostic=reference_diagnostic, | |
| pixel_metrics=identity_metrics(int(reference_u8.shape[0])), | |
| frames=reference_u8, | |
| ) | |
| for policy in policies: | |
| latent, diagnostic = generate_rollout( | |
| pipeline=pipeline, | |
| conditional_dict=conditional, | |
| seed=seed, | |
| device=device, | |
| policy=policy, | |
| ) | |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16): | |
| pixels = vae.decode_to_pixel(latent, use_cache=False) | |
| prediction_u8 = base.pixels_to_u8(pixels) | |
| metrics = base.frame_metrics( | |
| reference_u8=reference_u8, | |
| prediction_u8=prediction_u8, | |
| lpips_model=lpips_model, | |
| batch_size=4, | |
| device=device, | |
| ) | |
| write_result( | |
| output_root=output_root, | |
| mapping_row=mapping_row, | |
| policy=policy, | |
| seed=seed, | |
| diagnostic=diagnostic, | |
| pixel_metrics=compact_metrics(metrics), | |
| frames=prediction_u8, | |
| ) | |
| print( | |
| f"[result] global={global_index} strategy={policy.name} " | |
| f"F={diagnostic['full_calls']} R={diagnostic['reuse_calls']} " | |
| f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} " | |
| f"lpips={metrics['lpips']:.6f}", | |
| flush=True, | |
| ) | |
| if hasattr(vae.model, "clear_cache"): | |
| vae.model.clear_cache() | |
| del latent, pixels, prediction_u8, metrics | |
| torch.cuda.empty_cache() | |
| if hasattr(vae.model, "clear_cache"): | |
| vae.model.clear_cache() | |
| del conditional, reference_latent, reference_pixels, reference_u8 | |
| torch.cuda.empty_cache() | |
| def write_manifest( | |
| *, | |
| output_root: Path, | |
| gpu: str, | |
| shard_index: int, | |
| num_shards: int, | |
| rows: list[dict[str, Any]], | |
| policies: tuple[NaivePolicy, ...], | |
| seed: int, | |
| status: str, | |
| ) -> None: | |
| selected = "_".join(policy.name for policy in policies) | |
| selection_id = hashlib.sha256(selected.encode("utf-8")).hexdigest()[:8] | |
| manifest_name = f"naive_gpu{gpu}_shard{shard_index:02d}_{selection_id}.json" | |
| atomic_json( | |
| output_root / "generation_metrics/manifests" / manifest_name, | |
| { | |
| "status": status, | |
| "physical_gpu": gpu, | |
| "shard_index": shard_index, | |
| "num_shards": num_shards, | |
| "global_indices": [int(row["global_index"]) for row in rows], | |
| "num_prompts": len(rows), | |
| "strategies": [REFERENCE_POLICY.name, *(policy.name for policy in policies)], | |
| "seed": seed, | |
| }, | |
| ) | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--gpu", default=PHYSICAL_GPU) | |
| parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT) | |
| parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT) | |
| parser.add_argument("--shard-index", type=int, default=0) | |
| parser.add_argument("--num-shards", type=int, default=1) | |
| parser.add_argument( | |
| "--strategy", | |
| action="append", | |
| choices=EVALUATION_STRATEGY_NAMES, | |
| default=None, | |
| help="Baseline to generate; repeat to select several. Default: all six.", | |
| ) | |
| parser.add_argument( | |
| "--global-index", | |
| dest="global_indices", | |
| action="append", | |
| type=int, | |
| default=None, | |
| help="Process only this global mapping index; may be repeated.", | |
| ) | |
| parser.add_argument("--seed", type=int, default=0) | |
| parser.add_argument("--overwrite", action="store_true") | |
| return parser.parse_args() | |
| def main() -> None: | |
| args = parse_args() | |
| if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards: | |
| raise ValueError("Invalid shard index/number of shards") | |
| strategy_names = args.strategy or list(EVALUATION_STRATEGY_NAMES) | |
| if len(set(strategy_names)) != len(strategy_names): | |
| raise ValueError("--strategy values must be unique") | |
| policies = tuple(POLICIES_BY_NAME[name] for name in strategy_names) | |
| mapping = read_mapping(args.mapping.resolve()) | |
| if args.global_indices is None: | |
| rows = [ | |
| row | |
| for position, row in enumerate(mapping) | |
| if position % args.num_shards == args.shard_index | |
| ] | |
| else: | |
| requested = set(args.global_indices) | |
| mapping_by_global = {int(row["global_index"]): row for row in mapping} | |
| unknown = sorted(requested - set(mapping_by_global)) | |
| if unknown: | |
| raise ValueError(f"Unknown global indices: {unknown}") | |
| rows = [row for row in mapping if int(row["global_index"]) in requested] | |
| output_root = args.output_root.resolve() | |
| output_root.mkdir(parents=True, exist_ok=True) | |
| write_manifest( | |
| output_root=output_root, | |
| gpu=str(args.gpu), | |
| shard_index=args.shard_index, | |
| num_shards=args.num_shards, | |
| rows=rows, | |
| policies=policies, | |
| seed=args.seed, | |
| status="running", | |
| ) | |
| print( | |
| f"[setup] physical_gpu={args.gpu} shard={args.shard_index}/{args.num_shards} " | |
| f"prompts={len(rows)} strategies={','.join(strategy_names)}", | |
| flush=True, | |
| ) | |
| device = torch.device("cuda") | |
| torch.set_grad_enabled(False) | |
| vae, pipeline, text_encoder, lpips_model = load_models(device) | |
| completed = 0 | |
| try: | |
| for row in rows: | |
| if not args.overwrite and is_prompt_complete(output_root, row, policies): | |
| completed += 1 | |
| print( | |
| f"[cached] {completed}/{len(rows)} global={row['global_index']}", | |
| flush=True, | |
| ) | |
| continue | |
| run_prompt( | |
| mapping_row=row, | |
| policies=policies, | |
| output_root=output_root, | |
| seed=args.seed, | |
| vae=vae, | |
| pipeline=pipeline, | |
| text_encoder=text_encoder, | |
| lpips_model=lpips_model, | |
| device=device, | |
| ) | |
| completed += 1 | |
| print(f"[progress] {completed}/{len(rows)} prompts", flush=True) | |
| finally: | |
| write_manifest( | |
| output_root=output_root, | |
| gpu=str(args.gpu), | |
| shard_index=args.shard_index, | |
| num_shards=args.num_shards, | |
| rows=rows, | |
| policies=policies, | |
| seed=args.seed, | |
| status="complete" if completed == len(rows) else "failed", | |
| ) | |
| print(f"[complete] gpu={args.gpu} prompts={completed}/{len(rows)}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |