Spaces:
Running on Zero
Running on Zero
| """Configuration-only validation for TF/FCP workflows. | |
| No dataset, checkpoint, model, distributed group or CUDA is initialized here. | |
| The narrow auxiliary range check uses a CPU timestep table to preserve the scheduler's | |
| float32 endpoint semantics; the usual full range needs no tensor imports. | |
| """ | |
| import math | |
| import re | |
| # channels, backbone blocks, hidden width, FPS; these are family facts, not recipes. | |
| FAMILIES = { | |
| "Wan2.1-T2V-1.3B": (16, 30, 1536, 16), | |
| "Wan2.1-T2V-14B": (16, 40, 5120, 16), | |
| "Wan2.2-TI2V-5B": (48, 30, 3072, 24), | |
| } | |
| def fail(field, rule): | |
| raise ValueError(f"{field}: {rule}") | |
| def integer(value, field, minimum=None): | |
| if type(value) is not int or (minimum is not None and value < minimum): | |
| fail( | |
| field, | |
| f"must be a non-boolean integer >= {minimum}" | |
| if minimum is not None | |
| else "must be a non-boolean integer", | |
| ) | |
| return value | |
| def number(value, field, minimum=None, maximum=None, *, strict_min=False, strict_max=False): | |
| if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): | |
| fail(field, "must be a finite non-boolean number") | |
| if minimum is not None and (value < minimum or (strict_min and value == minimum)): | |
| fail(field, f"must be {'>' if strict_min else '>='} {minimum}") | |
| if maximum is not None and (value > maximum or (strict_max and value == maximum)): | |
| fail(field, f"must be {'<' if strict_max else '<='} {maximum}") | |
| return value | |
| def string(value, field, *, nullable=False, empty=False): | |
| if nullable and value is None: | |
| return | |
| if not isinstance(value, str) or (not empty and not value.strip()): | |
| fail(field, "must be a string" if empty else "must be a nonempty string") | |
| def choice(value, field, choices): | |
| if value not in choices: | |
| fail(field, f"must be one of {', '.join(map(str, choices))}") | |
| def items(value, field, *, nonempty=False): | |
| if not isinstance(value, list) or (nonempty and not value): | |
| fail(field, "must be a nonempty list" if nonempty else "must be a list") | |
| return value | |
| def increasing(values, field, *, minimum=0, nonempty=False): | |
| items(values, field, nonempty=nonempty) | |
| for value in values: | |
| integer(value, field, minimum) | |
| if any(left >= right for left, right in zip(values, values[1:], strict=False)): | |
| fail(field, "must be strictly increasing with no duplicates") | |
| def validate_sources(config, *, required): | |
| sources = config.get("data_sources") | |
| if sources is None and not required: | |
| return | |
| items(sources, "data_sources", nonempty=True) | |
| seen = set() | |
| for index, source in enumerate(sources): | |
| prefix = f"data_sources[{index}]" | |
| source_id = source.get("id", f"source_{index}") | |
| string(source_id, f"{prefix}.id") | |
| source_id = source_id.strip() | |
| if source_id in seen: | |
| fail(f"{prefix}.id", f"duplicate source ID {source_id!r}") | |
| seen.add(source_id) | |
| source["id"] = source_id | |
| for key in ("data_path", "annotation_path"): | |
| string(source.get(key), f"{prefix}.{key}") | |
| if not source["data_path"].rstrip("/\\").lower().endswith("shards"): | |
| fail(f"{prefix}.data_path", "must name a directory ending in 'shards'") | |
| number(source.get("weight"), f"{prefix}.weight", 0, strict_min=True) | |
| if required: | |
| for key in ("data_path", "annotation_path"): | |
| if config.get(key) is not None: | |
| fail(key, "must be null/omitted when using data_sources") | |
| def validate_auxiliary_range(bounds, shift): | |
| field = "future_prediction.aux_future_timestep_range" | |
| items(bounds, field) | |
| if len(bounds) != 2: | |
| fail(field, "requires two finite nonnegative endpoints") | |
| lower, upper = [number(value, field, 0) for value in bounds] | |
| if lower > upper: | |
| fail(field, "lower endpoint must not exceed upper endpoint") | |
| # Every retained training schedule starts at 1000. Avoid importing torch for | |
| # the reference interval; a narrower interval needs the exact CPU arithmetic. | |
| if lower <= 1000 <= upper: | |
| return | |
| import torch | |
| sigmas = torch.linspace(1.0, 0.0, 1001, dtype=torch.float32, device="cpu")[:-1] | |
| timesteps = (shift * sigmas / (1 + (shift - 1) * sigmas)) * 1000 | |
| if not ((timesteps >= lower) & (timesteps <= upper)).any().item(): | |
| fail(field, "closed interval selects no actual training scheduler timestep") | |
| def validate_future_prediction(config, family, chunks): | |
| fp = config["future_prediction"] | |
| prefix = "future_prediction." | |
| if type(fp["enabled"]) is not bool: | |
| fail(prefix + "enabled", "must be boolean") | |
| if fp["enabled"]: | |
| for field in ("future_offsets", "offset_loss_weights", "source_layers", "mini_dit_depth"): | |
| if field not in fp: | |
| fail( | |
| prefix + field, "is required for FCP; method changes do not construct a recipe" | |
| ) | |
| offsets = fp.get("future_offsets") | |
| if "future_offsets" in fp: | |
| increasing(offsets, prefix + "future_offsets", minimum=1, nonempty=True) | |
| if offsets[-1] >= chunks: | |
| fail(prefix + "future_offsets", f"every offset must be below {chunks} regular chunks") | |
| if "offset_loss_weights" in fp: | |
| weights = items(fp["offset_loss_weights"], prefix + "offset_loss_weights", nonempty=True) | |
| if offsets is not None and len(weights) != len(offsets): | |
| fail(prefix + "offset_loss_weights", "requires one weight per offset") | |
| for weight in weights: | |
| number(weight, prefix + "offset_loss_weights", 0) | |
| if not any(weight > 0 for weight in weights): | |
| fail(prefix + "offset_loss_weights", "at least one weight must be positive") | |
| layers = fp.get("source_layers") | |
| if "source_layers" in fp: | |
| increasing(layers, prefix + "source_layers", nonempty=True) | |
| if layers[-1] >= family[1]: | |
| fail(prefix + "source_layers", f"indices must be below {family[1]}") | |
| depth = None | |
| if "mini_dit_depth" in fp: | |
| depth = integer(fp["mini_dit_depth"], prefix + "mini_dit_depth", 1) | |
| mode = fp["init_mode"] | |
| choice(mode, prefix + "init_mode", ("copy", "random")) | |
| if mode == "random": | |
| if "init_source" in fp: | |
| fail( | |
| prefix + "init_source", "must be absent for random initialization (null is invalid)" | |
| ) | |
| elif fp["enabled"] or "init_source" in fp: | |
| if depth is None: | |
| fail(prefix + "mini_dit_depth", "is required with init_source") | |
| if "init_source" not in fp: | |
| if len(layers) != 1: | |
| fail(prefix + "init_source", "is required for multi-layer copy initialization") | |
| indices = list(range(layers[0] + 1, layers[0] + 1 + depth)) | |
| else: | |
| indices = items(fp["init_source"], prefix + "init_source", nonempty=True) | |
| if len(indices) == 1 and isinstance(indices[0], str): | |
| match = re.fullmatch(r"\s*(-?\d+)\s*:\s*(-?\d+)\s*", indices[0]) | |
| if not match: | |
| fail(prefix + "init_source", "range must use start:end (end exclusive)") | |
| start, end = map(int, match.groups()) | |
| if not 0 <= start < end <= family[1]: | |
| fail(prefix + "init_source", "range must lie inside backbone blocks") | |
| indices = list(range(start, end)) | |
| for index in indices: | |
| integer(index, prefix + "init_source", 0) | |
| if index >= family[1]: | |
| fail(prefix + "init_source", f"indices must be below {family[1]}") | |
| if len(indices) != depth or len(set(indices)) != len(indices): | |
| fail(prefix + "init_source", "distinct index count must equal mini_dit_depth") | |
| fp["init_source"] = indices | |
| fusion = fp["source_layer_fusion"] | |
| choice(fusion["mode"], prefix + "source_layer_fusion.mode", ("mlp",)) | |
| choice(fusion["activation"], prefix + "source_layer_fusion.activation", ("gelu",)) | |
| ratio = number( | |
| fusion["hidden_ratio"], prefix + "source_layer_fusion.hidden_ratio", 0, strict_min=True | |
| ) | |
| if not math.isfinite(ratio * family[2]) or round(ratio * family[2]) <= 0: | |
| fail( | |
| prefix + "source_layer_fusion.hidden_ratio", "must produce a positive finite MLP width" | |
| ) | |
| streams = fp["source_context"].get("streams") | |
| if streams is None: | |
| streams = ["noisy"] | |
| items(streams, prefix + "source_context.streams", nonempty=True) | |
| if any(not isinstance(value, str) or value not in ("clean", "noisy") for value in streams): | |
| fail(prefix + "source_context.streams", "only clean and noisy are supported") | |
| if len(set(streams)) != len(streams): | |
| fail(prefix + "source_context.streams", "duplicates are not allowed") | |
| fp["source_context"]["streams"] = [name for name in ("clean", "noisy") if name in streams] | |
| shift = fp.get("aux_future_timestep_shift", config["model_kwargs"]["timestep_shift"]) | |
| number(shift, prefix + "aux_future_timestep_shift", 0, strict_min=True) | |
| validate_auxiliary_range(fp["aux_future_timestep_range"], shift) | |
| def validate_config(config: dict, *, training: bool, world_size: int = 1) -> None: | |
| """Validate and normalize one final merged public config before side effects.""" | |
| for key in ( | |
| "image_or_video_shape", | |
| "num_frame_per_block", | |
| "ema_weight", | |
| "ema_start_step", | |
| "mixed_precision", | |
| "gradient_checkpointing", | |
| "sharding_strategy", | |
| "log_iters", | |
| "num_workers", | |
| "guidance_scale", | |
| ): | |
| if key not in config: | |
| fail(key, "is required; select a complete training recipe") | |
| bool_fields = """load_diffusion_pretrained generator_fsdp_sync_module_states | |
| independent_first_frame condition_only_first_frame i2v mixed_precision gradient_checkpointing | |
| separate_predictive_grad_clip save_resume_checkpoint require_raw_generator_checkpoint | |
| no_save no_visualize disable_wandb use_ema use_train_data""".split() | |
| for key in bool_fields: | |
| if type(config[key]) is not bool: | |
| fail(key, "must be boolean") | |
| model = config["model_name"] | |
| if not isinstance(model, str) or model not in FAMILIES: | |
| fail("model_name", f"must be one of {', '.join(FAMILIES)}") | |
| family = FAMILIES[model] | |
| positive_ints = """batch_size val_batch_size log_iters gc_interval | |
| timestep_bucket_log_interval num_frame_per_block sampling_steps | |
| long_inference_threshold vae_decode_chunk""".split() | |
| for key in positive_ints: | |
| integer(config[key], key, 1) | |
| for key in ( | |
| "num_workers", | |
| "val_num_workers", | |
| "lr_warmup_steps", | |
| "ema_start_step", | |
| "wds_sample_shuffle", | |
| "wds_sample_shuffle_initial", | |
| ): | |
| integer(config[key], key, 0) | |
| for key in ("max_iters", "snapshot_iters", "num_output_frames"): | |
| value = integer(config[key], key) | |
| if value != -1 and value <= 0: | |
| fail(key, "must be positive or -1") | |
| integer(config["test_sample_limit"], "test_sample_limit", -1) | |
| if config["max_checkpoints"] is not None: | |
| integer(config["max_checkpoints"], "max_checkpoints") | |
| seed = integer(config["seed"], "seed", 0) | |
| integer(world_size, "world_size", 1) | |
| if seed > 2**32 - 1 or (training and (seed == 0 or seed + world_size - 1 > 2**32 - 1)): | |
| fail("seed", "must fit NumPy uint32 for every rank; training also requires nonzero") | |
| string(config["optimizer"], "optimizer") | |
| config["optimizer"] = config["optimizer"].strip().lower() | |
| choice(config["optimizer"], "optimizer", ("adamw",)) | |
| choice(config["ema_checkpoint_mode"], "ema_checkpoint_mode", ("default",)) | |
| choice(config["sample_solver"], "sample_solver", ("unipc", "dpm++")) | |
| choice( | |
| config["sharding_strategy"], | |
| "sharding_strategy", | |
| ("full", "hybrid_full", "hybrid_zero2", "no_shard"), | |
| ) | |
| for key in ("generator_fsdp_wrap_strategy", "text_encoder_fsdp_wrap_strategy"): | |
| if key in config: | |
| choice(config[key], key, ("size",)) | |
| adamw = config["optimizer_config"]["adamw"] | |
| for key in ("lr", "weight_decay"): | |
| number(adamw[key], f"optimizer_config.adamw.{key}", 0) | |
| for key in ("beta1", "beta2"): | |
| number(adamw[key], f"optimizer_config.adamw.{key}", 0, 1, strict_max=True) | |
| increasing(config["lr_drop_steps"], "lr_drop_steps", minimum=1) | |
| if any(step <= config["lr_warmup_steps"] for step in config["lr_drop_steps"]): | |
| fail("lr_drop_steps", "each node must be after lr_warmup_steps") | |
| number(config["lr_drop_gamma"], "lr_drop_gamma", 0, 1, strict_min=True, strict_max=True) | |
| number(config["ema_weight"], "ema_weight", 0, 1, strict_max=True) | |
| number(config["caption_dropout_p"], "caption_dropout_p", 0, 1) | |
| number(config["guidance_scale"], "guidance_scale", 0) | |
| number(config["max_grad_norm"], "max_grad_norm", 0, strict_min=True) | |
| number(config["timestep_shift"], "timestep_shift", 0, strict_min=True) | |
| number( | |
| config["model_kwargs"]["timestep_shift"], "model_kwargs.timestep_shift", 0, strict_min=True | |
| ) | |
| string(config["negative_prompt"], "negative_prompt", empty=True) | |
| for key in ("wandb_host", "wandb_project"): | |
| string(config[key], key) | |
| for key in ( | |
| "wandb_entity", | |
| "wandb_save_dir", | |
| "expdir", | |
| "checkpoint_path", | |
| "generator_ckpt", | |
| "resume_checkpoint_path", | |
| "data_path", | |
| "annotation_path", | |
| "output_path", | |
| "checkpoint_name", | |
| "inference_train_data_source_id", | |
| ): | |
| string(config[key], key, nullable=True) | |
| shape = items(config["image_or_video_shape"], "image_or_video_shape") | |
| if len(shape) != 5: | |
| fail("image_or_video_shape", "must be [B,T,C,H,W]") | |
| for value in shape: | |
| integer(value, "image_or_video_shape", 1) | |
| _, frames, channels, height, width = shape | |
| if channels != family[0]: | |
| fail("image_or_video_shape", f"{model} requires C={family[0]}") | |
| if height % 2 or width % 2: | |
| fail("image_or_video_shape", "H and W must be divisible by the spatial patch size 2") | |
| independent = config["independent_first_frame"] | |
| if config["condition_only_first_frame"] and not independent: | |
| fail("condition_only_first_frame", "requires independent_first_frame=true") | |
| regular = frames - int(independent) | |
| block = config["num_frame_per_block"] | |
| if regular <= 0 or regular % block: | |
| fail("image_or_video_shape", "T must contain positive whole regular chunks after C0") | |
| fps = config.get("fps", family[3]) | |
| number(fps, "fps", 0, strict_min=True) | |
| if int(fps) != fps: | |
| fail("fps", "must be integer-valued") | |
| kwargs = config["model_kwargs"] | |
| window = integer(kwargs["local_attn_size"], "model_kwargs.local_attn_size") | |
| if window == -1: | |
| window = frames | |
| elif window <= 0: | |
| fail("model_kwargs.local_attn_size", "must be positive or -1 (automatic)") | |
| sink = integer(kwargs["sink_size"], "model_kwargs.sink_size", 0) | |
| if sink + block > window: | |
| fail("model_kwargs.local_attn_size", "must fit sink_size plus an incoming regular chunk") | |
| validate_future_prediction(config, family, regular // block) | |
| predictive = config["future_prediction"]["enabled"] | |
| separate = config["separate_predictive_grad_clip"] | |
| if separate and not predictive: | |
| fail("separate_predictive_grad_clip", "requires future_prediction.enabled=true") | |
| if separate: | |
| number(config["max_predictive_grad_norm"], "max_predictive_grad_norm", 0, strict_min=True) | |
| elif config["max_predictive_grad_norm"] is not None: | |
| fail("max_predictive_grad_norm", "must be null/absent when separate clipping is off") | |
| if config["save_resume_checkpoint"] and config["no_save"]: | |
| fail("save_resume_checkpoint", "cannot be enabled with no_save=true") | |
| single_image = config.get("image_path") is not None | |
| if training: | |
| for key in ( | |
| "image_path", | |
| "video_path", | |
| "prompt", | |
| "prompt_file", | |
| "v2v", | |
| "v2v_context_latents", | |
| ): | |
| if key in config: | |
| fail(key, "is an inference-only input") | |
| else: | |
| string(config.get("image_path"), "image_path", nullable=True) | |
| string(config.get("prompt"), "prompt", nullable=True) | |
| string(config.get("prompt_file"), "prompt_file", nullable=True) | |
| string(config.get("video_path"), "video_path", nullable=True) | |
| single_video = config.get("video_path") is not None | |
| if single_image and single_video: | |
| fail("video_path", "cannot combine with image_path") | |
| if type(config["v2v"]) is not bool: | |
| fail("v2v", "must be boolean") | |
| if single_video: | |
| config["v2v"] = True | |
| v2v = config["v2v"] | |
| if v2v and config["i2v"]: | |
| fail("v2v", "cannot combine with i2v") | |
| if v2v and single_image: | |
| fail("v2v", "cannot combine with image_path") | |
| if v2v and not (single_video or config["use_train_data"]): | |
| fail("v2v", "requires video_path or use_train_data=true") | |
| if single_video: | |
| if config["data_path"] is not None or config["use_train_data"]: | |
| fail("video_path", "cannot combine with data_path or use_train_data") | |
| if config["inference_train_data_source_id"] is not None: | |
| fail("video_path", "cannot combine with train_data_source_id") | |
| if not v2v and config["v2v_context_latents"] is not None: | |
| fail("v2v_context_latents", "requires v2v=true or video_path") | |
| inline_prompt = config.get("prompt") is not None | |
| file_prompt = config.get("prompt_file") is not None | |
| if inline_prompt and file_prompt: | |
| fail("prompt/prompt_file", "are mutually exclusive; choose only one") | |
| if file_prompt and not config["prompt_file"].lower().endswith(".txt"): | |
| fail("prompt_file", "must name a .txt file") | |
| if (single_image or single_video) != (inline_prompt or file_prompt): | |
| fail( | |
| "image_path/--prompt", | |
| "provide both --image_path or --video_path and either --prompt or --prompt_file", | |
| ) | |
| if single_image: | |
| if config["data_path"] is not None: | |
| fail("image_path", "cannot combine with data_path; select one input route") | |
| if config["use_train_data"] or config["inference_train_data_source_id"] is not None: | |
| fail("image_path", "cannot combine with use_train_data or train_data_source_id") | |
| # A single supplied first frame selects TI2V without a redundant --i2v flag. | |
| config["i2v"] = True | |
| validate_sources(config, required=training or config["use_train_data"]) | |
| duration = config["duration"] | |
| string(duration, "duration") | |
| if config["num_output_frames"] == -1 and not re.fullmatch(r"[1-9][0-9]*[smhSMH]", duration): | |
| fail("duration", "must be a positive whole duration with s/m/h units") | |
| if training: | |
| return | |
| if not config["i2v"] and not config["v2v"]: | |
| fail("i2v", "ordinary inference requires text plus supplied C0; enable i2v") | |
| if not config["v2v"] and not independent and block != 1: | |
| fail( | |
| "num_frame_per_block", | |
| "single-C0 inference without an independent first frame requires 1", | |
| ) | |
| if config["v2v"]: | |
| from src.utils.video_prefix import prefix_geometry | |
| if single_video and config["v2v_context_latents"] is None: | |
| # Metadata-dependent tail selection happens before GPU/model loading. | |
| minimum = 1 + block if independent else block | |
| prefix_geometry({**config, "v2v_context_latents": minimum}) | |
| else: | |
| prefix_geometry(config) | |
| source_id = config["inference_train_data_source_id"] | |
| if config["use_train_data"]: | |
| sources = config["data_sources"] | |
| if source_id is None and len(sources) == 1: | |
| source_id = sources[0]["id"] | |
| if source_id not in {source["id"] for source in sources}: | |
| fail("inference_train_data_source_id", "select one configured source for inference") | |
| config["inference_train_data_source_id"] = source_id | |
| else: | |
| if source_id is not None: | |
| fail("inference_train_data_source_id", "requires use_train_data=true") | |
| if not single_image and not single_video: | |
| string(config["data_path"], "data_path") | |
| if world_size != 1: | |
| fail("i2v", "external-image input requires non-distributed Python execution") | |