diff --git a/.gitignore b/.gitignore index 3353a8bdb..126a777d9 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,8 @@ local/ output experiment /data/* +!primus/backends/diffusion/data/ +!primus/backends/diffusion/data/** +primus/backends/diffusion/data/**/__pycache__/ +primus/backends/diffusion/data/**/*.pyc pp_simulation_result diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 2df6093f7..5134b6068 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -1,12 +1,13 @@ -# Wan Examples +# Diffusion Examples -Wan examples exercise the independent PyTorch Diffusion backend under -`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal -configs, see `primus/backends/diffusion/README.md`. +These examples exercise the independent PyTorch diffusion backend under +`primus/backends/diffusion`, including Wan video training and FLUX text-to-image +training. For backend details, data/checkpoint layout, and minimal configs, see +`primus/backends/diffusion/README.md`. ## Data -The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face: +Wan examples use the `zirui3/tiny-video-samples` smoke-test dataset: ```bash huggingface-cli download zirui3/tiny-video-samples \ @@ -22,6 +23,15 @@ Expected layout: data/*.mp4 ``` +FLUX examples use either: + +- precomputed encodings: a Hugging Face `datasets` directory with + `t5_encodings`, `clip_encodings`, `mean`, and `logvar` fields. +- raw image-text data: a local webdataset directory, a local JSONL file, a local + Hugging Face dataset, or a Hugging Face dataset repo. + +See `primus/backends/diffusion/README.md` for checkpoint and data layout details. + ## Run Set the shared `torchrun` environment first: @@ -34,7 +44,7 @@ export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} ``` -### Pretrain +### Wan Pretrain ```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ @@ -53,7 +63,45 @@ torchrun \ Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism when the model head count supports it. -### Posttrain +### FLUX Precomputed Pretrain + +```bash +DATASET_PATH=/data/flux_precomputed \ +EMPTY_ENCODINGS_PATH=/data/flux_empty_encodings \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +``` + +Set `PRETRAINED_PATH=/path/to/flux1-dev.safetensors` only when you want to +initialize the DiT from existing FLUX weights. Empty `PRETRAINED_PATH` means +random initialization. + +### FLUX Raw Image-Text Pretrain + +```bash +DATASET=cc12m-test \ +DATASET_PATH=/data/cc12m_test \ +T5_ENCODER=google/t5-v1_1-xxl \ +CLIP_ENCODER=openai/clip-vit-large-patch14 \ +VAE_CHECKPOINT=black-forest-labs/FLUX.1-dev/ae.safetensors \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml +``` + +For raw training, `T5_ENCODER`, `CLIP_ENCODER`, and `VAE_CHECKPOINT` can be local +paths or Hugging Face identifiers. + +### Wan Posttrain ```bash INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ @@ -70,5 +118,5 @@ torchrun \ The MI355X configs use Primus-style override sections such as `training`, `data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion -adapter normalizes those sections into the Wan model/dataset/trainer +adapter normalizes those sections into diffusion model/dataset/trainer arguments at runtime. diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml new file mode 100644 index 000000000..5310f2dc3 --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml @@ -0,0 +1,52 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_dev_t2i.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_dev_t2i-pretrain + + data: + dataset_path: ${DATASET_PATH:} + empty_encodings_path: ${EMPTY_ENCODINGS_PATH:} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml new file mode 100644 index 000000000..e3689a489 --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml @@ -0,0 +1,63 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-raw-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_dev_t2i.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + model: + config: + encoder: + t5_encoder: ${T5_ENCODER:google/t5-v1_1-xxl} + clip_encoder: ${CLIP_ENCODER:openai/clip-vit-large-patch14} + autoencoder: ${VAE_CHECKPOINT:black-forest-labs/FLUX.1-dev/ae.safetensors} + max_t5_length: ${MAX_T5_LENGTH:256} + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-raw-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_dev_t2i-raw-pretrain + + data: + dataset_type: raw + dataset: ${DATASET:cc12m-test} + dataset_format: ${DATASET_FORMAT:webdataset} + dataset_path: ${DATASET_PATH:} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + skip_low_resolution: false + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 9b7db9704..951eee590 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -11,9 +11,10 @@ so Wan training is a first-class part of the Primus `diffusion` backend. Supported scope: -- Model implementation: `wan` for Wan2.1 and Wan2.2. +- Model implementation: `wan` for Wan2.1/Wan2.2 and `flux` for FLUX.1-dev T2I. - Trainer: FSDP2 only. -- Sequence parallelism: Ulysses SP via `trainer.args.sp_size`. +- Sequence parallelism: Ulysses SP via `trainer.args.sp_size` for Wan. FLUX + currently requires `sp_size=1`. Wan-specific dependencies are kept out of top-level Primus requirements. See `runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt`. @@ -52,6 +53,135 @@ data: video_backend: imageio ``` +### FLUX Checkpoints And Data + +FLUX supports two training-data modes: + +- `dataset_type: precomputed`: samples already contain T5 encodings, CLIP + encodings, and VAE latent statistics. This is the fastest path and matches the + MLPerf/TorchTitan preprocessed recipe. +- `dataset_type: raw`: samples contain original image-text pairs. Primus loads + frozen T5, CLIP, and FLUX autoencoder modules and computes the encodings online + during training. This is useful for bring-up and custom datasets, but it is + slower and uses more memory. + +#### FLUX Pretrain Checkpoints + +For FLUX pretraining, the DiT checkpoint is optional: + +```yaml +model: + config: + load_from_pretrained_path: "" # empty means random DiT init +``` + +Raw image-text training needs frozen encoders/autoencoder. These can be local +paths or Hugging Face identifiers in the YAML: + +```yaml +model: + config: + encoder: + t5_encoder: google/t5-v1_1-xxl + clip_encoder: openai/clip-vit-large-patch14 + autoencoder: black-forest-labs/FLUX.1-dev/ae.safetensors +``` + +`t5_encoder` and `clip_encoder` are loaded with Hugging Face `from_pretrained`. +`autoencoder` accepts either a local safetensors path or `repo_id/filename`; +for example `black-forest-labs/FLUX.1-dev/ae.safetensors`. + +The shipped precomputed example uses random DiT initialization unless +`PRETRAINED_PATH` points at FLUX DiT weights. It expects preprocessed data and, +when prompt dropout is enabled, empty prompt encodings: + +```text +PRETRAINED_PATH=/path/to/flux1-dev.safetensors # optional +DATASET_PATH=/path/to/cc12m_preprocessed +EMPTY_ENCODINGS_PATH=/path/to/empty_encodings +``` + +The example config is: + +```text +examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +``` + +#### FLUX Raw Data + +Raw FLUX data is selected from YAML. The supported built-in dataset names mirror +TorchTitan's FLUX recipe: + +```yaml +data: + dataset_type: raw + dataset: cc12m-test # small smoke-test webdataset +``` + +For full CC12M webdataset training, use: + +```yaml +data: + dataset_type: raw + dataset: cc12m-wds # maps to Hugging Face dataset pixparse/cc12m-wds +``` + +Do not use full `cc12m-wds` as the default smoke test; it is large. For +`cc12m-test`, pass a local webdataset directory explicitly: + +```text +DATASET=cc12m-test +DATASET_PATH=/path/to/cc12m_test +``` + +The raw FLUX example loads frozen encoders online. The defaults use Hugging Face +IDs for T5/CLIP/AE, but each value can also be a local path: + +```text +DATASET=cc12m-test +DATASET_PATH=/path/to/cc12m_test +T5_ENCODER=google/t5-v1_1-xxl +CLIP_ENCODER=openai/clip-vit-large-patch14 +VAE_CHECKPOINT=black-forest-labs/FLUX.1-dev/ae.safetensors +``` + +The raw example config is +`examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml`. + +Raw FLUX data can also be a local JSONL file: + +```jsonl +{"image": "000001.jpg", "prompt": "a caption"} +{"image": "000002.jpg", "txt": "another caption"} +``` + +Use: + +```yaml +data: + dataset_type: raw + dataset_format: jsonl + dataset_path: /path/to/meta.jsonl + data_folder: /path/to/images + img_size: 256 +``` + +You can also point `dataset_path` at a local webdataset directory or a Hugging +Face dataset repo, depending on `dataset_format`: + +```yaml +data: + dataset_type: raw + dataset_format: hf_repo + dataset_path: pixparse/cc12m-wds +``` + +In short, the intended design is: checkpoint and dataset references live in YAML; +each value can be a local path when assets are staged manually, or a Hugging Face +identifier when `transformers`, `huggingface_hub`, or `datasets` should resolve +it. This keeps runs reproducible while still allowing local cache/download +behavior. + ## Checkpoints Each Wan model is a single Hugging Face repo that already bundles everything Wan @@ -115,6 +245,14 @@ before launching distributed training: python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +# FLUX precomputed config; DiT checkpoint is optional +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml + +# FLUX raw image-text config; requires T5/CLIP access and FLUX AE checkpoint +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml + # post-train config (validates modules.post_trainer) python3 runner/helpers/hooks/train/posttrain/diffusion/prepare.py \ --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml @@ -126,9 +264,9 @@ for debugging. ## Launch Training runs through the Primus CLI (`primus.cli.main train `) -under `torchrun`. Single-node and multi-node share **one** launch command: the -same `torchrun` invocation runs on every node, parameterized by the standard -rendezvous variables. The defaults below give a single-node 8-GPU run. +under `torchrun`. Single-node and multi-node share one launch shape; choose the +config path for the workload you want to run. The defaults below give a +single-node 8-GPU FLUX precomputed run. ```bash # distributed knobs (defaults = single node, 8 GPUs) @@ -138,25 +276,32 @@ export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export MASTER_PORT=${MASTER_PORT:-29500} export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +# FLUX precomputed. Set DATASET_PATH and EMPTY_ENCODINGS_PATH first. +export CONFIG=examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml + torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ --nproc_per_node="$GPUS_PER_NODE" \ - -m primus.cli.main train pretrain --config /path/to/wan_config.yaml + -m primus.cli.main train pretrain --config "$CONFIG" ``` - **Single node**: run as-is (the defaults above). - **Multi-node**: run the same command on each node with a shared `MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct `NODE_RANK` per node. World size is `NNODES * GPUS_PER_NODE`. -- **Post-train**: identical command with `train posttrain` and a posttrain config. +- **FLUX raw image-text**: set `CONFIG` to + `examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml` and set + `DATASET_PATH`, `T5_ENCODER`, `CLIP_ENCODER`, and `VAE_CHECKPOINT` as needed. +- **Wan/post-train**: use the same command with the matching Wan config; for + post-train use `train posttrain` and a posttrain config. Useful runtime knobs: -- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan training; use `sdpa` as the portable fallback or baseline. -- `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the - model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but - not `sp_size=8`. +- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` on ROCm. Use + `sdpa` as the portable fallback when AITER flash attention is unavailable. +- `trainer.args.sp_size`: Ulysses sequence parallel size for Wan. FLUX currently + requires `sp_size=1`. - `trainer.args.dp_replicate`: data parallel replication size. - `FIXED_TIMESTEP` and `FIXED_SEED`: optional debug variables for reproducible loss-alignment checks. @@ -173,9 +318,11 @@ export AMD_COMGR_CACHE_DIR=/path/to/large/cache/comgr ## Primus-Style Configs New examples should use Primus-style override sections and let the diffusion -adapter normalize them into Wan args. See: +adapter normalize them into diffusion model/dataset/trainer args. See: ```text +examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml +examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index b6b444ceb..ca352a003 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -13,8 +13,8 @@ from primus.core.utils.yaml_utils import nested_namespace_to_dict -class WanArgBuilder: - """Build the compact config object consumed by the Wan trainer.""" +class DiffusionArgBuilder: + """Build the compact config object consumed by diffusion trainers.""" DEFAULT_DATASET: dict[str, Any] = { "name": "wan", @@ -45,7 +45,7 @@ class WanArgBuilder: }, }, } - DEFAULT_TRAINER: dict[str, Any] = { + DEFAULT_WAN_TRAINER: dict[str, Any] = { "name": "fsdp2", "args": { "output_dir": "./output/wan", @@ -86,6 +86,65 @@ class WanArgBuilder: }, }, } + DEFAULT_FLUX_DATASET: dict[str, Any] = { + "name": "flux", + "config": { + "dataset_type": "precomputed", + "dataset_format": "hf_dataset", + "dataset_path": "/path/to/flux_precomputed_dataset", + "dataset": None, + "shuffle": True, + "processor_config": { + "processor_name": "flux_precomputed", + "processor_type": "flux_precomputed", + "prompt_dropout_prob": 0.0, + "empty_encodings_path": None, + "img_size": 256, + "skip_low_resolution": True, + }, + }, + } + DEFAULT_FLUX_TRAINER: dict[str, Any] = { + "name": "fsdp2", + "args": { + "output_dir": "./output/flux", + "per_device_train_batch_size": 1, + "per_device_eval_batch_size": 1, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": False, + "attention_backend": "flash_attn_aiter", + "learning_rate": 2.0e-4, + "lr_scheduler_type": "constant", + "warmup_steps": 0, + "weight_decay": 0.01, + "num_train_epochs": 1, + "max_steps": 100, + "logging_steps": 1, + "save_steps": 0, + "dataloader_num_workers": 4, + "report_to": "none", + "run_name": "flux-fsdp2", + "bf16": True, + "seed": 10007, + "optim": "adamw_torch", + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + "max_grad_norm": 1.0, + "fsdp2_wrap_target": "dit", + "fsdp_transformer_layer_cls_to_wrap": "DoubleStreamBlock,SingleStreamBlock", + "fsdp2_reshard_after_forward": True, + "save_strategy": "dit_only", + "sp_size": 1, + "dp_replicate": 1, + "flow_match_scheduler": { + "shift": 1, + "sigma_min": 0.0, + "extra_one_step": False, + "num_train_timesteps": 1000, + }, + }, + } def __init__(self) -> None: self._params: dict[str, Any] = {} @@ -96,16 +155,18 @@ def update(self, params: Any) -> None: elif isinstance(params, dict): self._params = copy.deepcopy(params) else: - raise TypeError(f"WanArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + raise TypeError( + f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}" + ) def finalize(self) -> SimpleNamespace: params = copy.deepcopy(self._params) if "model" not in params: - raise ValueError("Wan backend config requires a model preset.") + raise ValueError("Diffusion backend config requires a model preset.") for legacy_section in ("dataset", "trainer"): if legacy_section in params: raise ValueError( - f"Wan backend no longer accepts public `{legacy_section}` overrides. " + f"Diffusion backend no longer accepts public `{legacy_section}` overrides. " "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, " "`runtime`, and `metrics` sections instead." ) @@ -138,18 +199,48 @@ def _get_any(source: dict[str, Any], *keys: str) -> Any: return source[key] return None + @staticmethod + def _coerce_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + return bool(value) + + def _model_family(self, params: dict[str, Any]) -> str: + model = params.get("model") or {} + if not isinstance(model, dict): + raise TypeError(f"Diffusion model preset must be a dict, got {type(model).__name__}") + name = str(model.get("name", "")).strip().lower() + if not name: + raise ValueError("Diffusion model preset requires `model.name`.") + return name + + def _defaults_for_model(self, model_name: str) -> tuple[dict[str, Any], dict[str, Any]]: + if model_name == "wan": + return self.DEFAULT_DATASET, self.DEFAULT_WAN_TRAINER + if model_name == "flux": + return self.DEFAULT_FLUX_DATASET, self.DEFAULT_FLUX_TRAINER + raise ValueError(f"Unsupported diffusion model name: {model_name!r}") + def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]: - """Translate concise Primus-style Wan sections into trainer arguments. + """Translate concise Primus-style sections into trainer arguments. Public configs use high-level sections such as `training`, `data`, - `parallelism`, and `runtime`; internal defaults supply the compact - `dataset` and `trainer` objects consumed by the Wan runtime. + `parallelism`, and `runtime`; internal defaults supply compact + `dataset` and `trainer` objects consumed by the diffusion runtime. """ + model_name = self._model_family(params) + default_dataset, default_trainer = self._defaults_for_model(model_name) normalized = { "model": params["model"], - "dataset": copy.deepcopy(self.DEFAULT_DATASET), - "trainer": copy.deepcopy(self.DEFAULT_TRAINER), + "dataset": copy.deepcopy(default_dataset), + "trainer": copy.deepcopy(default_trainer), "stage": params.get("stage", "pretrain"), "primus": params.get("primus", {}), } @@ -171,10 +262,12 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("gradient_accumulation_steps",): ("gradient_accumulation_steps",), ("output_dir",): ("output_dir",), ("save_steps",): ("save_steps",), + ("save_strategy",): ("save_strategy",), ("run_name",): ("run_name",), ("num_train_epochs",): ("num_train_epochs",), ("dataloader_num_workers",): ("dataloader_num_workers",), ("resume_from_checkpoint",): ("resume_from_checkpoint",), + ("dataloader_num_workers",): ("dataloader_num_workers",), } for source_path, target_path in training_map.items(): value = self._get_any(training, *source_path) @@ -189,6 +282,14 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("text_tokenizer",): ("processor_config", "text_tokenizer"), ("processor_name",): ("processor_config", "processor_name"), ("processor_type",): ("processor_config", "processor_type"), + ("dataset_format",): ("dataset_format",), + ("dataset_type",): ("dataset_type",), + ("dataset",): ("dataset",), + ("shuffle",): ("shuffle",), + ("empty_encodings_path",): ("processor_config", "empty_encodings_path"), + ("prompt_dropout_prob",): ("processor_config", "prompt_dropout_prob"), + ("img_size",): ("processor_config", "img_size"), + ("skip_low_resolution",): ("processor_config", "skip_low_resolution"), } for source_path, target_path in data_map.items(): value = self._get_any(data, *source_path) @@ -237,11 +338,21 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("attention_backend",): ("attention_backend",), ("report_to",): ("report_to",), ("seed",): ("seed",), + ("bf16",): ("bf16",), + ("fp16",): ("fp16",), + ("gradient_checkpointing",): ("gradient_checkpointing",), ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",), } for source_path, target_path in runtime_map.items(): value = self._get_any(runtime, *source_path) if value is not None: + if target_path in { + ("bf16",), + ("fp16",), + ("gradient_checkpointing",), + ("fsdp2_reshard_after_forward",), + }: + value = self._coerce_bool(value) self._set_nested(trainer_args, target_path, value) log_freq = metrics.get("log_freq") @@ -259,4 +370,11 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, if resume_from_checkpoint is not None: self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint) + if model_name == "flux" and int(trainer_args.get("sp_size", 1)) != 1: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + return normalized + + +# Backwards-compatible import path for existing Wan-only callers. +WanArgBuilder = DiffusionArgBuilder diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py index 18ae84901..b72452684 100644 --- a/primus/backends/diffusion/data/__init__.py +++ b/primus/backends/diffusion/data/__init__.py @@ -4,4 +4,24 @@ # See LICENSE for license information. ############################################################################### -"""Data builders for the diffusion backend.""" +"""Data modules for the Primus diffusion backend.""" + +from .config import DatasetConfig +from .dataset import WanVideoDataset +from .flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from .processor import WanVideoDataProcessor + +__all__ = [ + "DatasetConfig", + "FluxPrecomputedDataset", + "FluxPrecomputedProcessor", + "FluxRawImageTextDataset", + "FluxRawImageTextProcessor", + "WanVideoDataProcessor", + "WanVideoDataset", +] diff --git a/primus/backends/diffusion/data/collator.py b/primus/backends/diffusion/data/collator.py new file mode 100644 index 000000000..c1fdb9e09 --- /dev/null +++ b/primus/backends/diffusion/data/collator.py @@ -0,0 +1,90 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import collections +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass +class VisionCollator: + processor: Any + + def pad_sequence(self, input_ids, batch_first, padding_value): + if self.processor.tokenizer.padding_side == "left": + input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids] + input_ids = torch.nn.utils.rnn.pad_sequence( + input_ids, batch_first=batch_first, padding_value=padding_value + ) + if self.processor.tokenizer.padding_side == "left": + input_ids = torch.flip(input_ids, [1]) + return input_ids + + def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + inputs = collections.defaultdict(list) + for instance in instances: + for key, values in instance.items(): + inputs[key].append(values) + + batched_inputs = {} + if "input_ids" in inputs: + input_ids = inputs.pop("input_ids") + input_ids = self.pad_sequence( + input_ids, + batch_first=True, + padding_value=self.processor.tokenizer.pad_token_id, + ) + batched_inputs["input_ids"] = input_ids + if "labels" in inputs: + labels = inputs.pop("labels") + labels = self.pad_sequence( + labels, + batch_first=True, + padding_value=-100, + ) + batched_inputs["labels"] = labels + + if "attention_mask" in inputs: + inputs.pop("attention_mask") + + attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id).long() + batched_inputs["attention_mask"] = attention_mask + + # for the other keys + for key, values in inputs.items(): + # Handle scalar/boolean values ( use_audio_in_video) + if isinstance(values[0], bool) or ( + isinstance(values[0], (int, float)) and not isinstance(values[0], torch.Tensor) + ): + batched_inputs[key] = values[0] + else: + batched_inputs[key] = torch.stack(values, dim=0) + return batched_inputs + + @property + def image_token_id(self): + return self.processor.tokenizer.convert_tokens_to_ids(self.processor.image_token) + + +@dataclass +class RawBatchCollator: + """ + A minimal collator that returns raw samples as a list of dicts. + + This is useful when model-specific padding/encoding happens in a separate + batch preparation step (e.g. processor.prepare_batch / trainer.prepare_batch), + keeping the DataLoader and trainer model-agnostic. + """ + + def __call__(self, instances: Sequence[dict]) -> list[dict]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + return list(instances) diff --git a/primus/backends/diffusion/data/config.py b/primus/backends/diffusion/data/config.py new file mode 100644 index 000000000..e425149ba --- /dev/null +++ b/primus/backends/diffusion/data/config.py @@ -0,0 +1,89 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from typing import Any, Literal + +from pydantic import BaseModel, field_validator + + +class Args(BaseModel): + extra_kwargs: dict[str, Any] = {} + + def to_dict(self): + return self.model_dump() + + def to_json(self): + return self.model_dump_json() + + +class ProcessorConfig(Args): + processor_name: str + processor_type: str + + +class DatasetConfig(Args): + dataset_type: str + data_folder: str + dataset_format: Literal["json", "jsonl", "csv", "yaml", "hf_dataset", "arrow"] + processor_config: dict | ProcessorConfig + + # Dataset configuration + dataset_path: str | None = None # Optional - used for external files + datasets: list[dict] | None = None # Optional - used for inline YAML definitions + shuffle: bool = True + data_seed: int | None = 42 + eval_dataset_path: str | None = None + + # Object storage configuration + object_storage: Literal["azure", "gcs", "none"] | None = "none" + bucket_name: str | None = None + + # Packing configuration + packing: bool | None = False + packing_strategy: str | None = None + packing_length: int | None = 32000 + filter_overlong: bool | None = True + filter_overlong_workers: int | None = 8 + max_length: int | None = None + + # Video configuration + video_sampling_strategy: Literal["fps", "frame_num"] | None = "fps" + video_max_pixels: int | None = 768 * 28 * 28 + video_max_frames: int | None = 768 + video_min_pixels: int | None = 3136 + frame_num: int | None = 64 + fps: int | None = 1 + video_backend: Literal["decord", "qwen_vl_utils", "qwen_omni_utils", "imageio"] | None = "qwen_vl_utils" + + @field_validator( + "video_max_pixels", + "video_max_frames", + "frame_num", + "fps", + "packing_length", + "max_length", + "filter_overlong_workers", + ) + @classmethod + def validate_positive_values(cls, v, info): + """Validate that numeric video and packing parameters are positive.""" + if v is not None and v <= 0: + field_name = info.field_name + raise ValueError(f"{field_name} must be positive, got {v}") + return v + + @field_validator("video_backend") + @classmethod + def validate_video_backend_migration(cls, v): + """Provide migration warning for deprecated torchvision backend.""" + if v == "torchvision": + raise ValueError( + "The 'torchvision' video backend has been removed. " + "Please use 'decord', 'qwen_vl_utils', or 'qwen_omni_utils' instead. " + "Migration guide: If you were using torchvision, 'decord' provides " + "similar functionality with better performance." + ) + return v diff --git a/primus/backends/diffusion/data/dataset.py b/primus/backends/diffusion/data/dataset.py new file mode 100644 index 000000000..139209923 --- /dev/null +++ b/primus/backends/diffusion/data/dataset.py @@ -0,0 +1,296 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Dataset class for loading video data from JSONL/CSV files. +""" + +import json +import os +from io import BytesIO +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from primus.backends.diffusion.utils.data_utils import smart_nframes +from primus.backends.diffusion.utils.vision_process import strip_file_uri + +from .collator import RawBatchCollator + + +def _get_decord_vr( + video_path: str, + *, + num_threads: int, +): + """Construct a decord.VideoReader.""" + from decord import VideoReader, cpu + + return VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + + +class BaseDataset(Dataset): + """Base dataset with minimal shared interface.""" + + def __init__(self, config, **kwargs) -> None: + super().__init__() + self.config = config + self.samples = [] + + def __len__(self): + return len(self.samples) + + +class WanVideoDataset(BaseDataset): + """Dataset for WanVideo training from JSONL/CSV.""" + + def __init__(self, processor, config={}): + """ + Initialize WanVideo dataset. + + Args: + processor: WanVideoDataProcessor instance + video_backend: Backend for video loading ('qwen_vl_utils' or 'decord') + """ + super().__init__(config) + self.config = config + self.data_path = Path(self.config.dataset_path) + self.processor = processor + + # Load metadata + self.samples = self._load_metadata() + self._sync_processor_video_limits() + + def _sync_processor_video_limits(self): + image_processor = None + if hasattr(self.processor, "processor") and hasattr(self.processor.processor, "image_processor"): + image_processor = self.processor.processor.image_processor + if image_processor is None: + return + if ( + getattr(image_processor, "max_pixels", None) is None + and getattr(self.config, "video_max_pixels", None) is not None + ): + image_processor.max_pixels = self.config.video_max_pixels + + def _load_metadata(self) -> list[dict]: + """Load metadata from JSONL or CSV file.""" + samples = [] + # TODO: add dummy data for debugging + if self.data_path.name == "dummy.jsonl": + return samples + + if self.data_path.suffix == ".jsonl": + with open(self.data_path) as f: + for line in f: + samples.append(json.loads(line.strip())) + elif self.data_path.suffix == ".json": + with open(self.data_path) as f: + samples = json.load(f) + elif self.data_path.suffix == ".csv": + import pandas as pd + + df = pd.read_csv(self.data_path) + samples = df.to_dict("records") + else: + raise ValueError(f"Unsupported file format: {self.data_path=}") + + return samples + + def _load_video_frames(self, video_path: str, data_folder=None, fps: int = 1) -> tuple[np.ndarray, float]: + """Load video frames using the specified backend.""" + if self.config.data_folder is not None: + video_path = os.path.join(self.config.data_folder, video_path) + + if self.config.video_backend == "decord": + return self.load_video_decord(video_path, fps) + elif self.config.video_backend == "qwen_vl_utils": + return self.load_video_qwen_vl_utils(video_path, fps) + elif self.config.video_backend == "imageio": + return self.load_video_imageio(video_path, fps) + else: + raise ValueError(f"Unsupported video backend: {self.config.video_backend}") + + def load_video_imageio(self, video_path, fps): + import imageio + + video_path = strip_file_uri(video_path) + reader = imageio.get_reader(video_path) + + # Sampling Strategy + total_frames = reader.count_frames() + total_frames = int(total_frames) + + if self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + + # DiffSynth Sequential Reading + frames = [] + for i, frame in enumerate(reader): + if i >= valid_nframes: + break + frames.append(frame) + + # Stack to numpy (T, H, W, C) + frames = np.array(frames) + sample_fps = fps # Simplification + + reader.close() + else: + reader.close() + raise NotImplementedError("Only frame_num strategy implemented for imageio backend") + + return frames, sample_fps + + def load_video_decord( + self, + video_path: str | list[str] | BytesIO, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Decord backend. + + Args: + video_path: Path to video file or BytesIO object + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + # Keep dataset logic simple: use a fixed thread count here. + from decord import VideoReader, cpu + + num_threads = 4 + align = os.getenv("ALIGN_WITH_DIFFSYNTH") == "1" + + if isinstance(video_path, BytesIO): + vr = VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + elif isinstance(video_path, list): + vr = _get_decord_vr(strip_file_uri(video_path[0]), num_threads=num_threads) + elif isinstance(video_path, str): + vr = _get_decord_vr(strip_file_uri(video_path), num_threads=num_threads) + else: + raise ValueError(f"Unsupported video path type: {type(video_path)}") + + total_frames, video_fps = len(vr), vr.get_avg_fps() + if self.config.video_sampling_strategy == "fps": + nframes = smart_nframes(total_frames, video_fps=video_fps, fps=fps) + # Maintain uniform sampling for FPS strategy + uniform_sampled_frames = np.linspace(0, total_frames - 1, nframes, dtype=int) + elif self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + + if align: + # sequential sampling align with diffsynth + uniform_sampled_frames = np.arange(valid_nframes, dtype=int) + else: + # uniform sampling + uniform_sampled_frames = np.linspace(0, total_frames - 1, valid_nframes, dtype=int) + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + frame_idx = uniform_sampled_frames.tolist() + spare_frames = vr.get_batch(frame_idx).asnumpy() + # spare_frames = torch.tensor(spare_frames).permute(0, 3, 1, 2) # Convert to TCHW format + + # Calculate sample_fps + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + # Return HWC numpy array to match processor expectations + return spare_frames, sample_fps # (frames, height, width, channels) + + def load_video_qwen_vl_utils( + self, + video_path: str, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Qwen VL utils. + + Args: + video_path: Path to video file + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + from primus.backends.diffusion.utils.vision_process import fetch_video + + video_dict = { + "type": "video", + "video": f"file://{video_path}", + "min_frames": 1, + "max_pixels": self.config.video_max_pixels, + "max_frames": self.config.video_max_frames, + "min_pixels": self.config.video_min_pixels, + } + + if self.config.video_sampling_strategy == "frame_num": + is_even = self.config.frame_num % 2 == 0 + n_frames = self.config.frame_num if is_even else self.config.frame_num + 1 + video_dict["nframes"] = n_frames + + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # Enforce VAE divisibility constraint + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + # else: keep 1 frame (or handle error) + + return frames, sample_fps + elif self.config.video_sampling_strategy == "fps": + video_dict["fps"] = fps + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # Also enforce for fps strategy + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + + return frames, sample_fps + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + """Get a single sample.""" + sample = self.samples[idx] + + # Load video frames + video_path = sample["video"] + video_frames, fps = self._load_video_frames(video_path) + + # Get prompt + prompt = sample.get("prompt", "") + # Return raw sample. + return { + "video_frames": video_frames, # np.ndarray, typically (T, H, W, C) + "prompt": prompt, + "fps": fps, + "num_frames": int(getattr(self.config, "frame_num", 0) or 0), + "video_path": str(video_path), + } + + def get_collator(self): + # Prefer raw collation; model-specific processing should happen in processor.prepare_batch. + return RawBatchCollator() diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py new file mode 100644 index 000000000..7e54ebc2c --- /dev/null +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -0,0 +1,342 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import io +import json +import math +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +from PIL import Image, ImageFile +from torch.utils.data import Dataset + +from primus.backends.diffusion.data.collator import RawBatchCollator + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +def _tensor_from_serialized_numpy(value: bytes) -> torch.Tensor: + array = np.load(io.BytesIO(value)) + tensor = torch.from_numpy(array) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + + +def _to_tensor(value: Any) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (bytes, bytearray)): + return _tensor_from_serialized_numpy(bytes(value)) + if isinstance(value, np.ndarray): + tensor = torch.from_numpy(value) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + return torch.as_tensor(value) + + +class FluxPrecomputedDataset(Dataset): + """Map-style dataset for precomputed FLUX text and VAE encodings.""" + + required_fields = ("t5_encodings", "clip_encodings", "mean", "logvar") + + def __init__(self, dataset_path: str): + if not dataset_path: + raise ValueError("FLUX precomputed dataset requires `dataset_path`.") + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX precomputed dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset_path = dataset_path + self.dataset = load_from_disk(dataset_path) + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + sample = self.dataset[idx] + missing = [field for field in self.required_fields if field not in sample] + if missing: + raise KeyError(f"FLUX precomputed sample missing fields: {missing}") + return {field: _to_tensor(sample[field]) for field in self.required_fields} + + def get_collator(self): + return RawBatchCollator() + + +class FluxRawImageTextDataset(Dataset): + """Map-style raw image-text dataset for online FLUX encoding.""" + + def __init__( + self, + *, + dataset_path: str | None, + dataset_format: str = "webdataset", + dataset_name: str | None = None, + data_folder: str | None = None, + ): + dataset_name = dataset_name or None + dataset_path, dataset_format = self._resolve_dataset(dataset_name, dataset_path, dataset_format) + self.dataset_path = dataset_path + self.dataset_format = dataset_format + self.dataset_name = dataset_name + self.data_folder = data_folder + self._records: list[dict[str, Any]] | None = None + + if dataset_format == "jsonl": + path = Path(dataset_path) + if not path.is_file(): + raise FileNotFoundError(f"FLUX raw jsonl metadata not found: {dataset_path}") + self._records = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + self.dataset = None + elif dataset_format == "hf_dataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw HF dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset = load_from_disk(dataset_path) + elif dataset_format == "hf_repo": + from datasets import load_dataset + + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_format == "webdataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw webdataset directory not found: {dataset_path}") + from datasets import load_dataset + + self.dataset = load_dataset( + "webdataset", + split="train", + data_dir=dataset_path, + data_files={"train": "*.tar"}, + ) + else: + raise ValueError("FLUX raw dataset_format must be one of: jsonl, hf_dataset, hf_repo, webdataset") + + @staticmethod + def _resolve_dataset( + dataset_name: str | None, + dataset_path: str | None, + dataset_format: str, + ) -> tuple[str, str]: + if dataset_name == "cc12m-test": + if not dataset_path: + raise ValueError("FLUX raw dataset `cc12m-test` requires `dataset_path` to be set.") + return dataset_path, "webdataset" + if dataset_name == "cc12m-wds": + return dataset_path or "pixparse/cc12m-wds", "hf_repo" + if not dataset_path: + raise ValueError("FLUX raw dataset requires either `dataset` or `dataset_path`.") + return dataset_path, dataset_format + + def __len__(self) -> int: + if self._records is not None: + return len(self._records) + return len(self.dataset) + + @staticmethod + def _prompt_from_sample(sample: dict[str, Any]) -> str: + for key in ("txt", "caption", "prompt", "text"): + if key in sample: + value = sample[key] + if isinstance(value, list): + value = value[0] + return str(value) + raise KeyError("FLUX raw sample requires one of: txt, caption, prompt, text") + + def _image_from_jsonl_sample(self, sample: dict[str, Any]) -> Image.Image: + image_key = "image" if "image" in sample else "jpg" if "jpg" in sample else "png" + image_path = Path(str(sample[image_key])) + if self.data_folder and not image_path.is_absolute(): + image_path = Path(self.data_folder) / image_path + return Image.open(image_path).convert("RGB") + + @staticmethod + def _image_from_dataset_sample(sample: dict[str, Any]) -> Image.Image: + if "image" in sample: + image = sample["image"] + elif "jpg" in sample: + image = sample["jpg"] + elif "png" in sample: + image = sample["png"] + else: + raise KeyError("FLUX raw dataset sample requires image, jpg, or png") + if isinstance(image, Image.Image): + return image.convert("RGB") + if isinstance(image, (bytes, bytearray)): + return Image.open(io.BytesIO(image)).convert("RGB") + return image.convert("RGB") + + def __getitem__(self, idx: int) -> dict[str, Any]: + if self._records is not None: + sample = self._records[idx] + image = self._image_from_jsonl_sample(sample) + else: + sample = self.dataset[idx] + image = self._image_from_dataset_sample(sample) + return {"image": image, "prompt": self._prompt_from_sample(sample)} + + def get_collator(self): + return RawBatchCollator() + + +@dataclass +class FluxPrecomputedProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.empty_encodings_path = self.config.get("empty_encodings_path") + self.img_size = int(self.config.get("img_size", 256) or 256) + self._empty_t5: torch.Tensor | None = None + self._empty_clip: torch.Tensor | None = None + + def build(self): + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + + @staticmethod + def _normalize_empty_encoding(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim >= 1 and tensor.shape[0] == 1: + return tensor[0] + return tensor + + def _check_empty_encoding_shapes(self, t5_encodings: torch.Tensor, clip_encodings: torch.Tensor) -> None: + # The empty encodings are broadcast-assigned into the per-sample encodings + # for dropped prompts, so their trailing (non-batch) shape must match the + # dataset encodings exactly. Fail fast with a clear message instead of a + # cryptic in-place assignment shape error mid-training. + assert self._empty_t5 is not None and self._empty_clip is not None + if self._empty_t5.shape != t5_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty T5 encoding shape " + f"{tuple(self._empty_t5.shape)} does not match the per-sample T5 encoding shape " + f"{tuple(t5_encodings.shape[1:])}. Regenerate t5_empty.npy with a matching " + "sequence length / hidden size." + ) + if self._empty_clip.shape != clip_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty CLIP encoding shape " + f"{tuple(self._empty_clip.shape)} does not match the per-sample CLIP encoding shape " + f"{tuple(clip_encodings.shape[1:])}. Regenerate clip_empty.npy with a matching shape." + ) + + def _load_empty_encodings(self): + if self._empty_t5 is not None and self._empty_clip is not None: + return + if not self.empty_encodings_path: + raise ValueError("FLUX prompt dropout requires `data.empty_encodings_path`.") + t5_path = os.path.join(self.empty_encodings_path, "t5_empty.npy") + clip_path = os.path.join(self.empty_encodings_path, "clip_empty.npy") + if not os.path.isfile(t5_path) or not os.path.isfile(clip_path): + raise FileNotFoundError( + "FLUX empty encodings must contain t5_empty.npy and clip_empty.npy " + f"under {self.empty_encodings_path}" + ) + self._empty_t5 = self._normalize_empty_encoding(torch.from_numpy(np.load(t5_path))) + self._empty_clip = self._normalize_empty_encoding(torch.from_numpy(np.load(clip_path))) + + @staticmethod + def _collate_raw(batch: Any) -> dict[str, torch.Tensor]: + if isinstance(batch, dict): + return {key: _to_tensor(value) for key, value in batch.items()} + if not isinstance(batch, Sequence) or not batch: + raise ValueError(f"FLUX prepare_batch expected non-empty sequence, got {type(batch).__name__}") + + keys = ("t5_encodings", "clip_encodings", "mean", "logvar") + collated: dict[str, torch.Tensor] = {} + for key in keys: + collated[key] = torch.stack([_to_tensor(sample[key]) for sample in batch], dim=0) + return collated + + def prepare_batch( + self, *, batch: Any, device: torch.device, dtype: torch.dtype + ) -> dict[str, torch.Tensor]: + tensors = self._collate_raw(batch) + for key, value in tensors.items(): + tensors[key] = value.to(device=device, dtype=dtype, non_blocking=True) + + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + assert self._empty_t5 is not None and self._empty_clip is not None + self._check_empty_encoding_shapes(tensors["t5_encodings"], tensors["clip_encodings"]) + bsz = tensors["t5_encodings"].shape[0] + drop_mask = torch.rand((bsz,), device=device) < self.prompt_dropout_prob + if drop_mask.any(): + tensors["t5_encodings"][drop_mask] = self._empty_t5.to(device=device, dtype=dtype) + tensors["clip_encodings"][drop_mask] = self._empty_clip.to(device=device, dtype=dtype) + + return tensors + + +@dataclass +class FluxRawImageTextProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.img_size = int(self.config.get("img_size", 256) or 256) + self.skip_low_resolution = bool(self.config.get("skip_low_resolution", True)) + + def build(self): + return + + def _process_image(self, image: Image.Image) -> torch.Tensor | None: + width, height = image.size + if self.skip_low_resolution and (width < self.img_size or height < self.img_size): + return None + + if width == self.img_size and height == self.img_size: + resized = image + elif width >= height: + new_width, new_height = math.ceil(self.img_size / height * width), self.img_size + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + left = int(torch.randint(0, new_width - self.img_size + 1, (1,)).item()) + resized = image.crop((left, 0, left + self.img_size, self.img_size)) + else: + new_width, new_height = self.img_size, math.ceil(self.img_size / width * height) + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + lower = int(torch.randint(0, new_height - self.img_size + 1, (1,)).item()) + resized = image.crop((0, lower, self.img_size, lower + self.img_size)) + + if resized.mode != "RGB": + resized = resized.convert("RGB") + np_img = np.array(resized).transpose((2, 0, 1)) + return torch.tensor(np_img).float() / 255.0 * 2.0 - 1.0 + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + if isinstance(batch, dict): + batch = [batch] + if not isinstance(batch, Sequence) or not batch: + raise ValueError( + f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}" + ) + + images: list[torch.Tensor] = [] + prompts: list[str] = [] + for sample in batch: + image = self._process_image(sample["image"]) + if image is None: + continue + prompt = str(sample.get("prompt", "")) + if self.prompt_dropout_prob > 0.0 and torch.rand(1).item() < self.prompt_dropout_prob: + prompt = "" + images.append(image) + prompts.append(prompt) + + if not images: + raise ValueError("FLUX raw batch contained no usable images after preprocessing") + return { + "image": torch.stack(images, dim=0).to(device=device, dtype=dtype, non_blocking=True), + "prompts": prompts, + } diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py new file mode 100644 index 000000000..675f3baa4 --- /dev/null +++ b/primus/backends/diffusion/data/processing_wanvideo.py @@ -0,0 +1,420 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import os +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image +from transformers import AutoTokenizer +from transformers.image_processing_utils import BaseImageProcessor +from transformers.image_utils import ImageInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from transformers.utils import TensorType, logging + +logger = logging.get_logger(__name__) + + +class WanVideoImageProcessor(BaseImageProcessor): + """ + Image/Video processor for WanVideo models. + + Args: + do_resize: Whether to resize the image/video frames. + size: Target size for resizing. + do_center_crop: Whether to center crop. + crop_size: Size for center cropping. + do_normalize: Whether to normalize pixel values. + image_mean: Mean values for normalization. + image_std: Standard deviation values for normalization. + do_convert_rgb: Whether to convert to RGB. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + do_resize: bool = True, + size: Optional[Dict[str, int]] = None, + do_center_crop: bool = True, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: bool = True, + max_pixels: Optional[int] = None, + height_division_factor: int = 16, + width_division_factor: int = 16, + **kwargs, + ): + super().__init__(**kwargs) + + self.do_resize = do_resize + self.size = size + self.do_center_crop = do_center_crop + self.crop_size = crop_size + self.do_normalize = do_normalize + self.image_mean = image_mean or [0.5, 0.5, 0.5] + self.image_std = image_std or [0.5, 0.5, 0.5] + self.do_convert_rgb = do_convert_rgb + self.max_pixels = max_pixels + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + + def resize( + self, + image: np.ndarray, + size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Resize image or video frame.""" + from PIL import Image as PILImage + + image = PILImage.fromarray(image.astype(np.uint8)) + # TODO: Here, we align with DiffSynth's resize logic for debugging(may remove in the future) + if os.getenv("ALIGN_WITH_DIFFSYNTH") == "1": + # Match DiffSynth logic: scale based on max dimension ratio + width, height = image.size + target_width = size["width"] + target_height = size["height"] + + scale = max(target_width / width, target_height / height) + new_width = round(width * scale) + new_height = round(height * scale) + + from torchvision.transforms import InterpolationMode + from torchvision.transforms import functional as F + + # DiffSynth uses torchvision.transforms.resize with BILINEAR + # We must use F.resize to match exactly (antialias behavior etc) + image = F.resize(image, (new_height, new_width), interpolation=InterpolationMode.BILINEAR) + else: + image = image.resize((size["width"], size["height"]), PILImage.LANCZOS) + return np.array(image) + + def _extract_hw(self, image: Union[np.ndarray, Image.Image]) -> Tuple[int, int]: + if isinstance(image, Image.Image): + width, height = image.size + return height, width + if isinstance(image, np.ndarray): + if image.ndim == 3: + return image.shape[0], image.shape[1] + if image.ndim == 2: + return image.shape[0], image.shape[1] + raise ValueError("Unsupported image type for size extraction.") + + def _dynamic_target_size(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]: + height, width = self._extract_hw(image) + if self.max_pixels is not None and width * height > self.max_pixels: + scale = (width * height / self.max_pixels) ** 0.5 + height = int(height / scale) + width = int(width / scale) + height = height // self.height_division_factor * self.height_division_factor + width = width // self.width_division_factor * self.width_division_factor + height = max(self.height_division_factor, height) + width = max(self.width_division_factor, width) + return {"height": height, "width": width} + + def _resolve_sizes( + self, + image: Union[np.ndarray, Image.Image], + size: Optional[Dict[str, int]] = None, + crop_size: Optional[Dict[str, int]] = None, + ) -> Tuple[Dict[str, int], Dict[str, int]]: + size = size if size is not None else self.size + crop_size = crop_size if crop_size is not None else self.crop_size + if size is None and crop_size is None: + size = self._dynamic_target_size(image) + crop_size = dict(size) + elif size is None: + size = dict(crop_size) + elif crop_size is None: + crop_size = dict(size) + size = { + "height": self._ceil_to_factor(size["height"], self.height_division_factor), + "width": self._ceil_to_factor(size["width"], self.width_division_factor), + } + crop_size = { + "height": self._ceil_to_factor(crop_size["height"], self.height_division_factor), + "width": self._ceil_to_factor(crop_size["width"], self.width_division_factor), + } + return size, crop_size + + @staticmethod + def _ceil_to_factor(value: int, factor: int) -> int: + if value % factor == 0: + return value + return (value + factor - 1) // factor * factor + + def center_crop( + self, + image: np.ndarray, + crop_size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Center crop image or video frame.""" + h, w = image.shape[:2] + crop_h, crop_w = crop_size["height"], crop_size["width"] + + top = (h - crop_h) // 2 + left = (w - crop_w) // 2 + if image.ndim == 3: + return image[top : top + crop_h, left : left + crop_w, :] + else: + return image[top : top + crop_h, left : left + crop_w] + + def normalize( + self, + image: np.ndarray, + mean: List[float], + std: List[float], + **kwargs, + ) -> np.ndarray: + """Normalize image or video frame.""" + image = image.astype(np.float32) / 255.0 + mean = np.array(mean).reshape(1, 1, -1) + std = np.array(std).reshape(1, 1, -1) + return (image - mean) / std + + def preprocess( + self, + images: ImageInput, + do_resize: Optional[bool] = None, + size: Optional[Dict[str, int]] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: Optional[bool] = None, + num_frames: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Preprocess images or video frames. + + Args: + images: Input images or video frames. + return_tensors: Type of tensors to return ("pt" for PyTorch). + + Returns: + Dictionary with preprocessed pixel values. + """ + do_resize = do_resize if do_resize is not None else self.do_resize + size = size if size is not None else self.size + do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop + crop_size = crop_size if crop_size is not None else self.crop_size + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + image_mean = image_mean if image_mean is not None else self.image_mean + image_std = image_std if image_std is not None else self.image_std + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + + # Handle single image or list of images (video frames) + if not isinstance(images, list): + images = [images] + + processed_images = [] + for image in images: + if isinstance(image, Image.Image): + image = np.array(image) + + # Handle 4D tensor case (batch of video frames) + if isinstance(image, np.ndarray) and image.ndim == 4: + # Process each frame in the batch + batch_processed_frames = [] + for i in range(image.shape[0]): + frame = image[i] # Get single frame + if frame.ndim == 3 and (frame.shape[0] == 3 or frame.shape[0] == 1): + frame = np.transpose(frame, (1, 2, 0)) # (C, H, W) -> (H, W, C) + + # Convert to RGB if needed + if do_convert_rgb and frame.shape[-1] != 3: + if len(frame.shape) == 2: # Grayscale + frame = np.stack([frame] * 3, axis=-1) + elif frame.shape[-1] == 4: # RGBA + frame = frame[..., :3] + + size_, crop_size_ = self._resolve_sizes(frame, size=size, crop_size=crop_size) + if do_resize: + frame = self.resize(frame, size_) + if do_center_crop: + frame = self.center_crop(frame, crop_size_) + + # Normalize + if do_normalize: + frame = self.normalize(frame, image_mean, image_std) + + batch_processed_frames.append(frame) + + # Stack frames back into 4D tensor + processed_image = np.stack(batch_processed_frames, axis=0) + else: + # Handle single image case (3D or 2D) + # Convert to RGB if needed + if do_convert_rgb and image.shape[-1] != 3: + if len(image.shape) == 2: # Grayscale + image = np.stack([image] * 3, axis=-1) + elif image.shape[-1] == 4: # RGBA + image = image[..., :3] + + size_, crop_size_ = self._resolve_sizes(image, size=size, crop_size=crop_size) + if do_resize: + image = self.resize(image, size_) + if do_center_crop: + image = self.center_crop(image, crop_size_) + + # Normalize + if do_normalize: + image = self.normalize(image, image_mean, image_std) + + processed_image = image + + processed_images.append(processed_image) + + # Stack frames for video + processed_images = np.stack(processed_images, axis=0) # B, T, H, W, C + + # Temporal Handling (Interpolate or Truncate) + if num_frames is not None: + current_frames = processed_images.shape[1] + if current_frames > num_frames: + logger.info(f"Truncating video frames from {current_frames} to {num_frames}") + processed_images = processed_images[:, :num_frames, ...] + elif current_frames < num_frames: + logger.info(f"Interpolating video frames from {current_frames} to {num_frames}") + # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C) + # Permute to (B, C, T, H, W) for interpolate + vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3) + + # Interpolate + vid_tensor = torch.nn.functional.interpolate( + vid_tensor, + size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]), + mode="trilinear", + align_corners=False, + ) + + # Permute back to (B, T, H, W, C) and convert to numpy + processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy() + + # Convert to tensor if requested + if return_tensors == "pt": + processed_images = torch.from_numpy(processed_images) + # Rearrange to (B, C, T, H, W) for video (since input was B, T, H, W, C) + if processed_images.ndim == 5: + processed_images = processed_images.permute(0, 4, 1, 2, 3) + elif processed_images.ndim == 4: + # (T, H, W, C) -> (C, T, H, W) if it was list of frames + processed_images = processed_images.permute(3, 0, 1, 2) + # Add batch dimension + # processed_images = processed_images.unsqueeze(0) + + return {"pixel_values": processed_images} + + +class WanVideoProcessor: + """ + Processor for WanVideo models, combining image/video processing and text tokenization. + + Args: + image_processor: Image/video processor instance. + tokenizer: Text tokenizer instance. + """ + + attributes = ["tokenizer", "image_processor"] + valid_kwargs = [ + "chat_template", + ] + tokenizer_class = "AutoTokenizer" + image_processor_class = "WanVideoImageProcessor" + + def __init__(self, image_processor=None, tokenizer=None, max_text_length: Optional[int] = None, **kwargs): + if image_processor is None: + image_processor = WanVideoImageProcessor(**kwargs) + if tokenizer is None: + # Default to T5 tokenizer for text encoding + try: + print("[Debug-processor], loading default tokenizer: google/umt5-xxl") + tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") + except Exception: + logger.warning("Could not load default tokenizer, using None") + tokenizer = None + + self.image_processor = image_processor + self.tokenizer = tokenizer + self.max_text_length = max_text_length + + def __call__( + self, + text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + images: Optional[ImageInput] = None, + videos: Optional[ImageInput] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Process text and image/video inputs. + + Args: + text: Text input(s) to tokenize. + images: Image input(s) to process. + videos: Video frames to process. + return_tensors: Type of tensors to return. + + Returns: + Dictionary with processed inputs. + """ + if text is None and images is None and videos is None: + raise ValueError("You must provide either text, images, or videos.") + + data = {} + + # Process text + if text is not None and self.tokenizer is not None: + max_length = kwargs.pop("max_length", None) + if max_length is None: + max_length = self.max_text_length or getattr(self.tokenizer, "model_max_length", 256) + text_inputs = self.tokenizer( + text, + return_tensors=return_tensors, + padding=True, + truncation=True, + max_length=max_length, + **kwargs, + ) + data.update(text_inputs) + + # Process images or video + if images is not None or videos is not None: + image_inputs = self.image_processor( + images or videos, + return_tensors=return_tensors, + **kwargs, + ) + data.update(image_inputs) + + return data + + def batch_decode(self, *args, **kwargs): + """Delegate to tokenizer's batch_decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """Delegate to tokenizer's decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + """Get model input names from components.""" + tokenizer_input_names = self.tokenizer.model_input_names if self.tokenizer else [] + image_processor_input_names = self.image_processor.model_input_names + return list(set(tokenizer_input_names + image_processor_input_names)) diff --git a/primus/backends/diffusion/data/processor.py b/primus/backends/diffusion/data/processor.py new file mode 100644 index 000000000..6434d24a6 --- /dev/null +++ b/primus/backends/diffusion/data/processor.py @@ -0,0 +1,182 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Processes video and text data for training. +""" + +from collections.abc import Sequence +from typing import Any + +import torch + + +class WanVideoDataProcessor: + """Standalone data processor for WanVideo training.""" + + def __init__(self, config, model_id=None): + self.config = config + self.model_id = model_id + self.processor = None + self.tokenizer = None + + def apply_prompt_template(self, hf_messages: str) -> str: + """Apply prompt template for WanVideo.""" + # WanVideo uses direct prompts without special formatting. + # Keep backward compatibility: + # - old path passed `hf_messages` list[dict] + # - new path passes plain prompt string + if isinstance(hf_messages, str): + return hf_messages + try: + return hf_messages[0]["content"][1]["text"] + except Exception: + # Best-effort fallback + return str(hf_messages) + + def save_pretrained(self, save_directory: str): + pass + + def build(self): + """Initialize the processor and tokenizer.""" + if self.processor is not None and self.tokenizer is not None: + return + + from transformers import AutoTokenizer + + from primus.backends.diffusion.data.processing_wanvideo import ( + WanVideoProcessor as WanVideoModelProcessor, + ) + + wanvideo_kwargs = self.config.get("extra_kwargs", {}) + max_text_length = self.config.get("max_text_length") + if max_text_length is not None: + wanvideo_kwargs.setdefault("max_text_length", max_text_length) + + # Load tokenizer if specified + if self.config.get("text_tokenizer", None) is not None: + self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("text_tokenizer")) + else: + self.tokenizer = None + self.processor = WanVideoModelProcessor(**wanvideo_kwargs, tokenizer=self.tokenizer) + + if self.tokenizer is None: + self.tokenizer = self.processor.tokenizer + + def _normalize_raw_batch(self, batch: Any) -> tuple[list[str], list[Any], int | None]: + """ + Normalize dataloader output into: + - prompts: list[str] + - frames_list: list[Any] (each item is typically np.ndarray (T,H,W,C)) + - num_frames: Optional[int] (only if consistent across samples) + """ + if isinstance(batch, dict): + prompts = batch.get("prompt") + frames_list = batch.get("video_frames") + num_frames = batch.get("num_frames", None) + + if prompts is None or frames_list is None: + raise ValueError("prepare_batch(dict) requires keys: 'prompt' and 'video_frames'") + if isinstance(prompts, str): + prompts = [prompts] + if not isinstance(prompts, (list, tuple)): + raise TypeError(f"'prompt' must be str or list[str], got {type(prompts)}") + if not isinstance(frames_list, (list, tuple)): + raise TypeError(f"'video_frames' must be list, got {type(frames_list)}") + + if isinstance(num_frames, (list, tuple)): + nfs = [int(x) for x in num_frames if x] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + elif num_frames: + num_frames = int(num_frames) + else: + num_frames = None + + return list(prompts), list(frames_list), num_frames + + # RawBatchCollator returns list[dict] + if not isinstance(batch, (list, tuple)) or not batch: + raise ValueError(f"prepare_batch expected non-empty list/tuple, got {type(batch)}") + + prompts = [str(ex.get("prompt", "")) for ex in batch] + frames_list = [ex.get("video_frames") for ex in batch] + + # Prefer per-example num_frames if present and consistent; else None. + nfs = [ex.get("num_frames") for ex in batch if ex.get("num_frames")] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + num_frames = int(num_frames) if num_frames else None + + return prompts, frames_list, num_frames + + def _tokenize_prompts(self, prompts: Sequence[str]) -> dict[str, torch.Tensor]: + formatted_prompts = [self.apply_prompt_template(p) for p in prompts] + return self.tokenizer( + formatted_prompts, + return_tensors="pt", + padding=self.config.get("padding_strategy", "max_length"), + truncation=True, + max_length=self.config.get("max_text_length", 512), + ) + + def _preprocess_videos(self, frames_list: Sequence[Any], num_frames: int | None) -> torch.Tensor: + if any(v is None for v in frames_list): + raise ValueError("prepare_batch got None in 'video_frames'") + video_inputs = self.processor.image_processor.preprocess( + list(frames_list), + num_frames=num_frames, + return_tensors="pt", + ) + if "pixel_values" not in video_inputs: + raise KeyError("image_processor.preprocess must return dict with key 'pixel_values'") + pixel_values = video_inputs["pixel_values"] + if not isinstance(pixel_values, torch.Tensor): + raise TypeError(f"pixel_values must be torch.Tensor, got {type(pixel_values)}") + if pixel_values.ndim != 5: + raise ValueError(f"Expected pixel_values shape [B,C,T,H,W], got {tuple(pixel_values.shape)}") + return pixel_values + + def _assemble_model_batch( + self, *, pixel_values: torch.Tensor, text_inputs: dict[str, torch.Tensor] + ) -> dict[str, Any]: + if "input_ids" not in text_inputs or "attention_mask" not in text_inputs: + raise KeyError("tokenizer output must include 'input_ids' and 'attention_mask'") + + out: dict[str, Any] = { + "video": pixel_values, # [B,C,T,H,W] + "input_ids": text_inputs["input_ids"], + "attention_mask": text_inputs["attention_mask"], + } + + # Legacy convenience fields (some models expect these keys). + out.setdefault("num_frames", int(pixel_values.shape[2])) + out.setdefault("height", int(pixel_values.shape[3])) + out.setdefault("width", int(pixel_values.shape[4])) + return out + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + """ + Convert raw dataloader batch into model inputs. + + Expected raw batch from `RawBatchCollator`: + - batch: list[dict] with keys: + - video_frames: np.ndarray (T,H,W,C) or equivalent + - prompt: str + - num_frames: int (optional) + + Output matches current training pipelines: + - video: Tensor [B, C, T, H, W] + - input_ids: Tensor [B, L] + - attention_mask: Tensor [B, L] + - (optional) height/width/num_frames for legacy wan/wan_new paths + """ + if self.processor is None or self.tokenizer is None: + # Lazy init to keep behavior robust across call sites. + self.build() + + prompts, frames_list, num_frames = self._normalize_raw_batch(batch) + text_inputs = self._tokenize_prompts(prompts) + pixel_values = self._preprocess_videos(frames_list, num_frames=num_frames) + return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs) diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py index b5a3705a2..4573e0a65 100644 --- a/primus/backends/diffusion/data/registrations/__init__.py +++ b/primus/backends/diffusion/data/registrations/__init__.py @@ -4,4 +4,4 @@ # See LICENSE for license information. ############################################################################### -"""Registered diffusion dataset builders.""" +"""Dataset registrations.""" diff --git a/primus/backends/diffusion/data/registrations/flux.py b/primus/backends/diffusion/data/registrations/flux.py new file mode 100644 index 000000000..234868970 --- /dev/null +++ b/primus/backends/diffusion/data/registrations/flux.py @@ -0,0 +1,38 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.utils.log import logger + + +def build_flux_dataset(dataset_config: dict): + processor_config = dataset_config.get("processor_config", {}) or {} + dataset_type = str(dataset_config.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + processor = FluxRawImageTextProcessor(processor_config) + processor.build() + dataset = FluxRawImageTextDataset( + dataset_path=dataset_config.get("dataset_path"), + dataset_format=dataset_config.get("dataset_format", "webdataset"), + dataset_name=dataset_config.get("dataset"), + data_folder=dataset_config.get("data_folder"), + ) + logger.info(f"Built FLUX raw image-text dataset with {len(dataset)} samples") + elif dataset_type == "precomputed": + processor = FluxPrecomputedProcessor(processor_config) + processor.build() + dataset = FluxPrecomputedDataset(dataset_config["dataset_path"]) + logger.info(f"Built FLUX precomputed dataset with {len(dataset)} samples") + else: + raise ValueError("FLUX dataset_type must be either 'precomputed' or 'raw'") + return dataset, processor diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index 76492cb2b..f832f352a 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -4,173 +4,25 @@ # See LICENSE for license information. ############################################################################### -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from PIL import Image -from torch.utils.data import Dataset -from torchvision.transforms import functional as F -from transformers import AutoTokenizer - -from primus.backends.diffusion.utils.vision_process import fetch_video, strip_file_uri - - -class WanVideoProcessor: - """Tokenize prompts and normalize video tensors for Wan training.""" - - def __init__(self, config: dict[str, Any]): - self.config = config - self.max_text_length = int(config.get("max_text_length", 512)) - tokenizer_path = config.get("text_tokenizer") - if not tokenizer_path: - raise ValueError("Wan dataset processor requires `text_tokenizer`.") - trust_remote_code = bool(config.get("trust_remote_code", False)) - self.tokenizer = AutoTokenizer.from_pretrained( - tokenizer_path, - trust_remote_code=trust_remote_code, - ) - - extra_kwargs = config.get("extra_kwargs", {}) or {} - size = extra_kwargs.get("size", {}) or {} - self.height = int(size.get("height", 480)) - self.width = int(size.get("width", 832)) - self.image_mean = torch.tensor(extra_kwargs.get("image_mean", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - self.image_std = torch.tensor(extra_kwargs.get("image_std", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - - def tokenize(self, prompt: str) -> dict[str, torch.Tensor]: - encoded = self.tokenizer( - prompt, - max_length=self.max_text_length, - padding="max_length", - truncation=True, - return_tensors="pt", - ) - return { - "input_ids": encoded["input_ids"].squeeze(0).long(), - "attention_mask": encoded["attention_mask"].squeeze(0).long(), - } - - def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - if video_tchw.ndim != 4: - raise ValueError(f"Expected video tensor [T,C,H,W], got shape={tuple(video_tchw.shape)}") - video_tchw = video_tchw[:, :3].float() - if video_tchw.max() > 2: - video_tchw = video_tchw / 255.0 - video_tchw = F.resize(video_tchw, [self.height, self.width], antialias=True) - video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() - return (video_cthw - self.image_mean) / self.image_std - - def prepare_batch( - self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype - ) -> dict[str, Any]: - return batch - - -class WanVideoDataset(Dataset): - def __init__(self, config: dict[str, Any], processor: WanVideoProcessor): - self.config = config - self.processor = processor - self.dataset_path = Path(config.get("dataset_path", "")) - self.data_folder = Path(config.get("data_folder", "")) - self.frame_num = int(config.get("frame_num", 81)) - self.video_backend = str(config.get("video_backend", "imageio")).lower() - - if not self.dataset_path.exists(): - raise FileNotFoundError(f"Wan dataset metadata not found: {self.dataset_path}") - with self.dataset_path.open(encoding="utf-8") as f: - self.samples = [json.loads(line) for line in f if line.strip()] - if not self.samples: - raise ValueError(f"Wan dataset metadata is empty: {self.dataset_path}") - - def __len__(self) -> int: - return len(self.samples) - - def _resolve_video_path(self, value: str) -> str: - if value.startswith(("http://", "https://", "file://")): - return value - path = Path(value) - if not path.is_absolute() and self.data_folder: - path = self.data_folder / path - return str(path) - - def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - total_frames = int(video_tchw.shape[0]) - if total_frames <= 0: - raise ValueError("Video contains no frames.") - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long() - return video_tchw[idx] - - def _read_video_imageio(self, path: str) -> torch.Tensor: - import imageio.v3 as iio - - path = strip_file_uri(path) - frames = [] - for frame in iio.imiter(path): - image = Image.fromarray(frame).convert("RGB") - frames.append(torch.from_numpy(np.asarray(image).copy())) - if not frames: - raise ValueError(f"Video contains no frames: {path}") - return torch.stack(frames).permute(0, 3, 1, 2) - - def _read_video_decord(self, path: str) -> torch.Tensor: - import decord - - path = strip_file_uri(path) - vr = decord.VideoReader(path) - total_frames = len(vr) - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() - return torch.from_numpy(vr.get_batch(idx).asnumpy()).permute(0, 3, 1, 2) - - def _read_video(self, path: str) -> torch.Tensor: - if self.video_backend == "imageio": - video = self._sample_video(self._read_video_imageio(path)) - elif self.video_backend == "decord": - video = self._read_video_decord(path) - else: - video = fetch_video( - { - "video": path, - "nframes": self.frame_num, - "resized_height": self.processor.height, - "resized_width": self.processor.width, - } - ) - return self.processor.normalize_video(video) - - def __getitem__(self, index: int) -> dict[str, Any]: - sample = self.samples[index] - prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" - video_key = sample.get("video") or sample.get("video_path") - if not video_key: - raise KeyError(f"Wan dataset sample missing `video`: index={index}") - - item = self.processor.tokenize(str(prompt)) - item["video"] = self._read_video(self._resolve_video_path(str(video_key))) - if "seed" in sample: - item["seed"] = int(sample["seed"]) - return item - - @staticmethod - def get_collator(): - def collate(samples: list[dict[str, Any]]) -> dict[str, Any]: - batch = { - "video": torch.stack([sample["video"] for sample in samples]), - "input_ids": torch.stack([sample["input_ids"] for sample in samples]), - "attention_mask": torch.stack([sample["attention_mask"] for sample in samples]), - } - if any("seed" in sample for sample in samples): - batch["seed"] = torch.tensor([sample.get("seed", 0) for sample in samples], dtype=torch.long) - return batch - - return collate - - -def build_wan_dataset(config: dict[str, Any]): - processor = WanVideoProcessor(config.get("processor_config", {}) or {}) - dataset = WanVideoDataset(config, processor) +"""Register Wan dataset builder.""" + +from primus.backends.diffusion.data import ( + DatasetConfig, + WanVideoDataProcessor, + WanVideoDataset, +) +from primus.backends.diffusion.utils.log import logger + + +def build_wan_dataset(dataset_config: dict): + processor_config = dataset_config["processor_config"] + processor = WanVideoDataProcessor(processor_config) + processor.build() + logger.info("Built data processor") + + dataset = WanVideoDataset( + processor=processor, + config=DatasetConfig(**dataset_config), + ) + logger.info(f"Built dataset with {len(dataset)} samples") return dataset, processor diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 344c293b3..878ea3fea 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from primus.backends.diffusion.argument_builder import WanArgBuilder +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.core.utils.module_utils import log_rank_0 @@ -42,17 +42,18 @@ def setup_backend_path(self, backend_path=None) -> str: return resolved_str def convert_config(self, params: Any): - builder = WanArgBuilder() + builder = DiffusionArgBuilder() builder.update(params) - wan_args = builder.finalize() + diffusion_args = builder.finalize() # convert_config is also called by the standalone prepare hook, where the # Primus logger may not be initialized yet; guard the informational log. try: - log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") + model_name = getattr(getattr(diffusion_args, "model", {}), "get", lambda *_: None)("name") + log_rank_0(f"[Primus:DiffusionAdapter] Converted Primus module params -> {model_name} args") except Exception: # Standalone prepare hooks may run before the Primus logger is bound. pass - return wan_args + return diffusion_args def load_trainer_class(self, stage: str = "pretrain"): if stage in ("pretrain", "posttrain", "sft"): diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 1b51b6a4c..b1b7fd650 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -15,11 +15,11 @@ class DiffusionPretrainTrainer(BaseTrainer): - """Primus lifecycle wrapper for Wan diffusion training.""" + """Primus lifecycle wrapper for diffusion backend training.""" def __init__(self, backend_args: Any): super().__init__(backend_args=backend_args) - self.wan_trainer = None + self.diffusion_trainer = None @staticmethod def _as_dict(value: Any) -> dict: @@ -48,6 +48,8 @@ def setup(self): if importlib.util.find_spec(package) is None ] video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") + if dataset_cfg.get("name") == "flux" and importlib.util.find_spec("datasets") is None: + missing.append("datasets") if video_backend == "imageio" and importlib.util.find_spec("imageio") is None: missing.append("imageio") if video_backend == "decord" and importlib.util.find_spec("decord") is None: @@ -55,7 +57,7 @@ def setup(self): if missing: raise RuntimeError( "Diffusion backend missing required Python packages: " - f"{', '.join(missing)}. Install the Wan diffusion training extras first." + f"{', '.join(missing)}. Install the diffusion training extras first." ) if attention_backend: @@ -99,7 +101,7 @@ def init(self): ) model = get_model_builder(model_name)(model_config) dataset, processor = get_dataset_builder(dataset_name)(dataset_config) - self.wan_trainer = get_trainer_builder(trainer_name)( + self.diffusion_trainer = get_trainer_builder(trainer_name)( model=model, dataset=dataset, processor=processor, @@ -107,11 +109,11 @@ def init(self): ) def train(self): - if self.wan_trainer is None: + if self.diffusion_trainer is None: raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().") - self.wan_trainer.train() - self.wan_trainer.save_model() + self.diffusion_trainer.train() + self.diffusion_trainer.save_model() def cleanup(self, on_error: bool = False): try: diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py index 2798dc3fb..7503f134c 100644 --- a/primus/backends/diffusion/models/__init__.py +++ b/primus/backends/diffusion/models/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Wan model exports for the Primus Wan backend.""" +"""Model exports for the Primus diffusion backend.""" from .wan import WanForTraining diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py new file mode 100644 index 000000000..7f4567d8a --- /dev/null +++ b/primus/backends/diffusion/models/flux/__init__.py @@ -0,0 +1,23 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.model import ( + Flux, + FluxParams, + flux_1_dev_params, +) +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) + +__all__ = [ + "Flux", + "FluxForTraining", + "FluxFlowMatchTrainPipeline", + "FluxParams", + "flux_1_dev_params", +] diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py new file mode 100644 index 000000000..a19d28a24 --- /dev/null +++ b/primus/backends/diffusion/models/flux/adapter.py @@ -0,0 +1,112 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) +from primus.backends.diffusion.models.interface import GenAIModel + + +@dataclass +class FluxConfigShim: + raw: dict + + def save_pretrained(self, save_directory: str): + import json + import os + + os.makedirs(save_directory, exist_ok=True) + with open(os.path.join(save_directory, "flux_config.json"), "w") as f: + json.dump(self.raw, f, indent=2, sort_keys=True) + + def to_dict(self): + return self.raw + + +class FluxForTraining(GenAIModel, nn.Module): + """Thin Primus adapter around the FLUX DiT backbone.""" + + def __init__( + self, + *, + dit: nn.Module, + train_pipeline: FluxFlowMatchTrainPipeline, + model_config: Any, + autoencoder: nn.Module | None = None, + t5_encoder: nn.Module | None = None, + clip_encoder: nn.Module | None = None, + raw_config: dict | None = None, + trainable_modules: str | None = None, + ): + super().__init__() + self.dit = dit + self.autoencoder = autoencoder + self.t5_encoder = t5_encoder + self.clip_encoder = clip_encoder + self.train_pipeline = train_pipeline + self.model_config = model_config + self.trainable_modules = trainable_modules + self.config = FluxConfigShim(raw=raw_config or {}) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def freeze_except(self): + mode = ( + self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit" + ).lower() + + def freeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(False) + + def unfreeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(True) + + freeze(self) + if mode in ("dit", "diffusion", "backbone"): + unfreeze(self.dit) + elif mode == "all": + unfreeze(self) + else: + unfreeze(self.dit) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + if hasattr(self.dit, "gradient_checkpointing"): + self.dit.gradient_checkpointing = True + + def forward(self, *args, **kwargs): + if len(args) >= 1 and isinstance(args[0], dict): + return self.forward_train(args[0], scheduler=kwargs.get("scheduler")) + raise TypeError("FluxForTraining.forward expects (batch_dict, scheduler)") + + def forward_train(self, batch: dict[str, Any], scheduler: Any = None) -> dict[str, torch.Tensor]: + del scheduler + return self.train_pipeline.compute_loss( + dit=self.dit, + autoencoder=self.autoencoder, + t5_encoder=self.t5_encoder, + clip_encoder=self.clip_encoder, + batch=batch, + model_config=self.model_config, + ) + + def forward_inference(self, batch: dict[str, Any], **kwargs): + raise NotImplementedError("FLUX inference pipeline is not wired for Primus training yet") diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py new file mode 100644 index 000000000..9ff6f3e56 --- /dev/null +++ b/primus/backends/diffusion/models/flux/autoencoder.py @@ -0,0 +1,315 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import torch +from einops import rearrange +from safetensors.torch import load_file as load_sft +from torch import Tensor, nn + + +@dataclass +class AutoEncoderParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + out_ch: int = 3 + ch_mult: list[int] | None = None + num_res_blocks: int = 2 + z_channels: int = 16 + scale_factor: float = 0.3611 + shift_factor: float = 0.1159 + + def __post_init__(self): + if self.ch_mult is None: + self.ch_mult = [1, 2, 4, 4] + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + bsz, channels, height, width = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + return rearrange(h_, "b 1 (h w) c -> b c h w", h=height, w=width, c=channels, b=bsz) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: Tensor) -> Tensor: + h = self.conv1(swish(self.norm1(x))) + h = self.conv2(swish(self.norm2(h))) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: Tensor) -> Tensor: + return self.conv(nn.functional.pad(x, (0, 1, 0, 1), mode="constant", value=0)) + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv(x) + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + return self.conv_out(swish(self.norm_out(h))) + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = ch * ch_mult[self.num_resolutions - 1] + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + self.up.insert(0, up) + + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + upscale_dtype = next(self.up.parameters()).dtype + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h).to(upscale_dtype) + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + return self.conv_out(swish(self.norm_out(h))) + + +class DiagonalGaussian(nn.Module): + def __init__(self, sample: bool = True, chunk_dim: int = 1): + super().__init__() + self.sample = sample + self.chunk_dim = chunk_dim + + def forward(self, z: Tensor) -> Tensor: + mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim) + if not self.sample: + return mean + std = torch.exp(0.5 * logvar) + return mean + std * torch.randn_like(mean) + + +class AutoEncoder(nn.Module): + def __init__(self, params: AutoEncoderParams, sample_z: bool = True): + super().__init__() + self.params = params + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian(sample=sample_z) + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: Tensor) -> Tensor: + z = self.reg(self.encoder(x)) + return self.scale_factor * (z - self.shift_factor) + + def decode(self, z: Tensor) -> Tensor: + z = z / self.scale_factor + self.shift_factor + return self.decoder(z) + + def forward(self, x: Tensor) -> Tensor: + return self.decode(self.encode(x)) + + +def load_autoencoder( + ckpt_path: str | None, + params: AutoEncoderParams | None = None, + *, + dtype: torch.dtype = torch.bfloat16, + sample_z: bool = True, +) -> AutoEncoder: + ae = AutoEncoder(params or AutoEncoderParams(), sample_z=sample_z) + if ckpt_path: + ckpt_path = _resolve_checkpoint_path(ckpt_path, default_filename="ae.safetensors") + state = load_sft(ckpt_path, device="cpu") + missing, unexpected = ae.load_state_dict(state, strict=False) + if missing: + raise ValueError(f"FLUX autoencoder checkpoint missing {len(missing)} keys; first={missing[:3]}") + if unexpected: + raise ValueError( + f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}" + ) + return ae.to(dtype=dtype) + + +def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) -> str: + path = Path(path_or_repo_file).expanduser() + if path.exists(): + return str(path) + if path_or_repo_file.startswith(("/", "./", "../", "~")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") + + parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + elif len(parts) == 2: + repo_id = path_or_repo_file + filename = default_filename + else: + raise FileNotFoundError( + f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}" + ) + + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) diff --git a/primus/backends/diffusion/models/flux/conditioner.py b/primus/backends/diffusion/models/flux/conditioner.py new file mode 100644 index 000000000..d63779c7c --- /dev/null +++ b/primus/backends/diffusion/models/flux/conditioner.py @@ -0,0 +1,54 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer + + +class HFEmbedder(nn.Module): + def __init__(self, version: str, max_length: int, **hf_kwargs): + super().__init__() + self.is_clip = "clip" in version.lower() or version.startswith("openai") + self.max_length = max_length + self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" + + if self.is_clip: + self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs) + else: + self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs) + + self.hf_module = self.hf_module.eval().requires_grad_(False) + for param in self.hf_module.parameters(): + if not param.is_contiguous(): + param.data = param.data.contiguous() + + @property + def device(self) -> torch.device: + return next(self.hf_module.parameters()).device + + def forward(self, text: list[str]) -> Tensor: + batch_encoding = self.tokenizer( + text, + truncation=True, + max_length=self.max_length, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + outputs = self.hf_module( + input_ids=batch_encoding["input_ids"].to(self.device), + attention_mask=None, + output_hidden_states=False, + ) + return outputs[self.output_key] diff --git a/primus/backends/diffusion/models/flux/configuration_flux.py b/primus/backends/diffusion/models/flux/configuration_flux.py new file mode 100644 index 000000000..6808fcac7 --- /dev/null +++ b/primus/backends/diffusion/models/flux/configuration_flux.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from types import SimpleNamespace + + +class FluxTrainingConfig(SimpleNamespace): + def to_dict(self) -> dict: + return dict(self.__dict__) diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py new file mode 100644 index 000000000..475be057b --- /dev/null +++ b/primus/backends/diffusion/models/flux/layers.py @@ -0,0 +1,211 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.math import attention, rope + + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + n_axes = ids.shape[-1] + emb = torch.cat([rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3) + return emb.unsqueeze(1) + + +def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> Tensor: + t = time_factor * t + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale + + +class QKNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.norm = QKNorm(head_dim) + self.proj = nn.Linear(dim, dim) + + def forward(self, x: Tensor, pe: Tensor) -> Tensor: + qkv = self.qkv(x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + x = attention(q, k, v, pe=pe) + return self.proj(x) + + +@dataclass +class ModulationOut: + shift: Tensor + scale: Tensor + gate: Tensor + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) + return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None + + +class DoubleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): + super().__init__() + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.num_heads = num_heads + self.hidden_size = hidden_size + self.img_mod = Modulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + self.txt_mod = Modulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]: + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + img_modulated = (1 + img_mod1.scale) * self.img_norm1(img) + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + txt_modulated = (1 + txt_mod1.scale) * self.txt_norm1(txt) + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + q = torch.cat((txt_q, img_q), dim=2) + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + attn = attention(q, k, v, pe=pe) + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(nn.Module): + def __init__( + self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None + ): + super().__init__() + self.hidden_dim = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) + self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) + self.norm = QKNorm(head_dim) + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp_act = nn.GELU(approximate="tanh") + self.modulation = Modulation(hidden_size, double=False) + + def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + attn = attention(q, k, v, pe=pe) + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + return self.linear(x) diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py new file mode 100644 index 000000000..829d2f395 --- /dev/null +++ b/primus/backends/diffusion/models/flux/math.py @@ -0,0 +1,44 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from einops import rearrange +from torch import Tensor + +from primus.backends.diffusion.attention import attention as backend_attention + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: + q, k = apply_rope(q, k, pe) + x = backend_attention( + q=rearrange(q, "B H L D -> B L H D"), + k=rearrange(k, "B H L D -> B L H D"), + v=rearrange(v, "B H L D -> B L H D"), + dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16, + ) + return rearrange(x, "B L H D -> B L (H D)") + + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + if dim % 2 != 0: + raise ValueError(f"RoPE dimension must be even, got {dim}") + scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + return rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2).float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py new file mode 100644 index 000000000..e5ebcce5a --- /dev/null +++ b/primus/backends/diffusion/models/flux/model.py @@ -0,0 +1,161 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +import torch +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.layers import ( + DoubleStreamBlock, + EmbedND, + LastLayer, + MLPEmbedder, + SingleStreamBlock, + timestep_embedding, +) + + +@dataclass +class FluxParams: + in_channels: int + out_channels: int + vec_in_dim: int + context_in_dim: int + hidden_size: int + mlp_ratio: float + num_heads: int + depth: int + depth_single_blocks: int + axes_dim: list[int] + theta: int + qkv_bias: bool + guidance_embed: bool + + def to_dict(self) -> dict: + return asdict(self) + + +def flux_1_dev_params(**overrides) -> FluxParams: + values = { + "in_channels": 64, + "out_channels": 64, + "vec_in_dim": 768, + "context_in_dim": 4096, + "hidden_size": 3072, + "mlp_ratio": 4.0, + "num_heads": 24, + "depth": 19, + "depth_single_blocks": 38, + "axes_dim": [16, 56, 56], + "theta": 10000, + "qkv_bias": True, + "guidance_embed": True, + } + values.update(overrides) + return FluxParams(**values) + + +class Flux(nn.Module): + """Transformer model for flow matching on packed latent sequences.""" + + def __init__(self, params: FluxParams): + super().__init__() + self.params = params + self.in_channels = params.in_channels + self.out_channels = params.out_channels + if params.hidden_size % params.num_heads != 0: + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" + ) + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError(f"Got axes_dim={params.axes_dim}, expected sum={pe_dim}") + + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.gradient_checkpointing = False + self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim) + self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) + self.guidance_in = ( + MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + ) + self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + ) + for _ in range(params.depth) + ] + ) + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio) + for _ in range(params.depth_single_blocks) + ] + ) + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) + + def _checkpoint_double(self, block: nn.Module, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, txt, vec, pe, use_reentrant=False) + + def _checkpoint_single(self, block: nn.Module, img: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, vec, pe, use_reentrant=False) + + def forward( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor | None = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + img = self.img_in(img) + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("FLUX guidance-distilled models require a guidance tensor.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + use_checkpoint = self.training and self.gradient_checkpointing + for block in self.double_blocks: + if use_checkpoint: + img, txt = self._checkpoint_double(block, img, txt, vec, pe) + else: + img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + + img = torch.cat((txt, img), 1) + for block in self.single_blocks: + if use_checkpoint: + img = self._checkpoint_single(block, img, vec, pe) + else: + img = block(img, vec=vec, pe=pe) + img = img[:, txt.shape[1] :, ...] + return self.final_layer(img, vec) diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py new file mode 100644 index 000000000..4bea0e954 --- /dev/null +++ b/primus/backends/diffusion/models/flux/train_pipeline.py @@ -0,0 +1,180 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +from primus.backends.diffusion.models.flux.utils import ( + create_position_encoding_for_latents, + generate_latent_from_mean_logvar, + pack_latents, +) + + +@dataclass +class FluxFlowMatchTrainPipelineConfig: + autoencoder_scale_factor: float = 0.3611 + autoencoder_shift_factor: float = 0.1159 + guidance: float = 1.0 + + +class FluxFlowMatchTrainPipeline: + """Flow-matching loss for FLUX using precomputed or online encodings.""" + + def __init__(self, cfg: FluxFlowMatchTrainPipelineConfig | None = None): + self.cfg = cfg or FluxFlowMatchTrainPipelineConfig() + + @staticmethod + def _require_module(module: torch.nn.Module | None, name: str) -> torch.nn.Module: + if module is None: + raise ValueError(f"FLUX raw image-text training requires `{name}` to be configured.") + return module + + @staticmethod + def _align_module_dtype(module: torch.nn.Module, *, device: torch.device, dtype: torch.dtype) -> None: + """Move a frozen encoder to the target device/dtype only when it differs. + + `nn.Module.to(dtype=...)` casts floating-point params/buffers only, leaving + integer buffers (e.g. token position ids) untouched. We inspect the first + floating-point parameter to avoid re-casting on every step. + """ + current = next((p for p in module.parameters() if p.is_floating_point()), None) + if current is None: + module.to(device=device) + return + if current.dtype != dtype or current.device != device: + module.to(device=device, dtype=dtype) + + def _prepare_precomputed( + self, + *, + batch: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + model_config: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t5_encodings = batch["t5_encodings"].to(device=device, dtype=dtype, non_blocking=True) + clip_encodings = batch["clip_encodings"].to(device=device, dtype=dtype, non_blocking=True) + mean = batch["mean"].to(device=device, dtype=dtype, non_blocking=True) + logvar = batch["logvar"].to(device=device, dtype=dtype, non_blocking=True) + + latents = generate_latent_from_mean_logvar(mean, logvar) + scale = float(getattr(model_config, "autoencoder_scale_factor", self.cfg.autoencoder_scale_factor)) + shift = float(getattr(model_config, "autoencoder_shift_factor", self.cfg.autoencoder_shift_factor)) + labels = (latents - shift) * scale + return labels, t5_encodings, clip_encodings + + def _prepare_raw( + self, + *, + batch: dict[str, Any], + autoencoder: torch.nn.Module | None, + t5_encoder: torch.nn.Module | None, + clip_encoder: torch.nn.Module | None, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ae = self._require_module(autoencoder, "model.config.encoder.autoencoder") + t5 = self._require_module(t5_encoder, "model.config.encoder.t5_encoder") + clip = self._require_module(clip_encoder, "model.config.encoder.clip_encoder") + image = batch["image"].to(device=device, dtype=dtype, non_blocking=True) + prompts = batch.get("prompts") + if not isinstance(prompts, list): + raise ValueError("FLUX raw batch requires `prompts` as list[str].") + + # The frozen encoders are built in bf16 regardless of the training compute + # dtype; align them with the DiT dtype so running them on `image`/inputs of + # `dtype` does not raise a dtype-mismatch error (e.g. for fp32 runs). + self._align_module_dtype(ae, device=device, dtype=dtype) + self._align_module_dtype(t5, device=device, dtype=dtype) + self._align_module_dtype(clip, device=device, dtype=dtype) + ae.eval() + t5.eval() + clip.eval() + with torch.no_grad(): + labels = ae.encode(image).to(device=device, dtype=dtype) + t5_encodings = t5(prompts).to(device=device, dtype=dtype) + clip_encodings = clip(prompts).to(device=device, dtype=dtype) + return labels, t5_encodings, clip_encodings + + def compute_loss( + self, + *, + dit: torch.nn.Module, + batch: dict[str, Any], + model_config: Any, + autoencoder: torch.nn.Module | None = None, + t5_encoder: torch.nn.Module | None = None, + clip_encoder: torch.nn.Module | None = None, + ) -> dict[str, torch.Tensor]: + if batch.get("sp_group") is not None: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + + device = next(dit.parameters()).device + dtype = next(dit.parameters()).dtype + if "image" in batch: + labels, t5_encodings, clip_encodings = self._prepare_raw( + batch=batch, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + device=device, + dtype=dtype, + ) + else: + required = ("t5_encodings", "clip_encodings", "mean", "logvar") + missing = [key for key in required if key not in batch] + if missing: + raise ValueError(f"FLUX precomputed batch is missing required keys: {missing}") + labels, t5_encodings, clip_encodings = self._prepare_precomputed( + batch=batch, + device=device, + dtype=dtype, + model_config=model_config, + ) + + bsz = labels.shape[0] + noise = torch.randn_like(labels) + timesteps = torch.rand((bsz,), device=device, dtype=dtype) + sigmas = timesteps.view(-1, 1, 1, 1) + noisy_latents = (1 - sigmas) * labels + sigmas * noise + target = noise - labels + + _, _, latent_height, latent_width = noisy_latents.shape + # Position ids are integer grid indices consumed by RoPE; build them in + # float32 (independent of the model compute dtype) so that indices remain + # exactly representable. bf16 only represents integers up to 256 exactly, + # which would silently corrupt positions for larger latent grids. + img_ids = create_position_encoding_for_latents( + bsz, + latent_height, + latent_width, + position_dim=3, + device=device, + dtype=torch.float32, + ) + txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=torch.float32) + noisy_latents = pack_latents(noisy_latents) + target = pack_latents(target) + + guidance_value = float(getattr(model_config, "guidance", self.cfg.guidance)) + guidance = torch.full((bsz,), guidance_value, device=device, dtype=dtype) + pred = dit( + img=noisy_latents, + img_ids=img_ids, + txt=t5_encodings, + txt_ids=txt_ids, + y=clip_encodings, + timesteps=timesteps, + guidance=guidance, + ) + loss = F.mse_loss(pred.float(), target.float(), reduction="sum") / target.numel() + return {"loss": loss, "log_metrics": {"mse": loss.detach()}} diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py new file mode 100644 index 000000000..00d0a2acf --- /dev/null +++ b/primus/backends/diffusion/models/flux/utils.py @@ -0,0 +1,49 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import torch +from torch import Tensor + +PATCH_HEIGHT = 2 +PATCH_WIDTH = 2 +LATENT_CHANNELS = 16 +IMAGE_LATENT_SIZE_RATIO = 8 + + +def generate_latent_from_mean_logvar(mean: Tensor, logvar: Tensor) -> Tensor: + return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean) + + +def create_position_encoding_for_latents( + bsz: int, + latent_height: int, + latent_width: int, + position_dim: int = 3, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> Tensor: + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + position_encoding = torch.zeros(height, width, position_dim, device=device, dtype=dtype) + position_encoding[:, :, 1] = torch.arange(height, device=device, dtype=dtype).unsqueeze(1) + position_encoding[:, :, 2] = torch.arange(width, device=device, dtype=dtype).unsqueeze(0) + return position_encoding.view(1, height * width, position_dim).repeat(bsz, 1, 1) + + +def pack_latents(x: Tensor) -> Tensor: + bsz, channels, latent_height, latent_width = x.shape + if latent_height % PATCH_HEIGHT != 0 or latent_width % PATCH_WIDTH != 0: + raise ValueError( + "FLUX latents must have height and width divisible by 2, " f"got shape={tuple(x.shape)}" + ) + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + x = x.unfold(2, PATCH_HEIGHT, PATCH_HEIGHT).unfold(3, PATCH_WIDTH, PATCH_WIDTH) + x = x.permute(0, 2, 3, 1, 4, 5) + return x.reshape(bsz, height * width, channels * PATCH_HEIGHT * PATCH_WIDTH) diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py new file mode 100644 index 000000000..89e8f3e23 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/flux.py @@ -0,0 +1,212 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import glob +import os +from typing import Any + +import torch +from safetensors.torch import load_file as safe_load_file + +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.autoencoder import ( + AutoEncoderParams, + load_autoencoder, +) +from primus.backends.diffusion.models.flux.conditioner import HFEmbedder +from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig +from primus.backends.diffusion.models.flux.model import Flux, flux_1_dev_params +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, + FluxFlowMatchTrainPipelineConfig, +) +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import count_parameters + + +def _strip_known_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + prefixes = ("module.", "dit.", "model.") + out: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + stripped = key + changed = True + while changed: + changed = False + for prefix in prefixes: + if stripped.startswith(prefix): + stripped = stripped[len(prefix) :] + changed = True + out[stripped] = value + return out + + +def _load_state_dict(path: str) -> dict[str, torch.Tensor]: + if path.endswith(".safetensors"): + return dict(safe_load_file(path)) + obj = torch.load(path, map_location="cpu") + if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): + obj = obj["model"] + if not isinstance(obj, dict): + raise ValueError(f"Unsupported checkpoint format at {path}") + return obj + + +def _candidate_weight_files(path: str) -> list[str]: + if os.path.isfile(path): + return [path] + if not os.path.exists(path): + resolved = _resolve_hf_checkpoint(path, default_filename="flux1-dev.safetensors") + if resolved: + return [resolved] + candidates: list[str] = [] + for fname in ("flux1-dev.safetensors", "dit_model.safetensors", "model.safetensors"): + candidate = os.path.join(path, fname) + if os.path.exists(candidate): + candidates.append(candidate) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.safetensors"))) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.bin"))) + return candidates + + +def _resolve_hf_checkpoint(path_or_repo_file: str, *, default_filename: str) -> str | None: + if path_or_repo_file.startswith(("/", "./", "../", "~")): + return None + parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + return None + if len(parts) < 2: + return None + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + else: + repo_id = path_or_repo_file + filename = default_filename + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) + + +def _load_flux_weights(dit: torch.nn.Module, pretrained_path: str) -> None: + candidates = _candidate_weight_files(pretrained_path) + if not candidates: + raise FileNotFoundError(f"No FLUX DiT weights found under {pretrained_path}") + + merged: dict[str, torch.Tensor] = {} + for ckpt in candidates: + merged.update(_strip_known_prefixes(_load_state_dict(ckpt))) + + result = dit.load_state_dict(merged, strict=False) + logger.info( + "Loaded FLUX DiT weights. " + f"files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + + +def _build_flux_dit(params) -> Flux: + local_rank = os.environ.get("LOCAL_RANK") + use_cuda = local_rank is not None and torch.cuda.is_available() + device = torch.device(f"cuda:{local_rank}") if use_cuda else torch.device("cpu") + old_dtype = torch.get_default_dtype() + try: + if use_cuda: + torch.set_default_dtype(torch.bfloat16) + with torch.device(device): + dit = Flux(params) + finally: + torch.set_default_dtype(old_dtype) + return dit + + +def build_flux_model(model_config: dict[str, Any]): + """ + YAML shape: + model_config: + load_from_pretrained_path: /path/to/flux1-dev.safetensors + config: + model_variant: flux-dev + trainable_modules: dit + guidance: 1.0 + params: {... optional FluxParams overrides ...} + """ + cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {}) + variant = cfg_dict.get("model_variant", "flux-dev") + if variant != "flux-dev": + raise ValueError(f"Only FLUX.1-dev is currently supported, got model_variant={variant!r}") + + params_overrides = dict(cfg_dict.get("params", {}) or {}) + params = flux_1_dev_params(**params_overrides) + dit = _build_flux_dit(params) + + pretrained_path = model_config.get("load_from_pretrained_path") or model_config.get("pretrained_path") + if pretrained_path: + logger.info(f"Loading FLUX DiT weights from {pretrained_path}") + _load_flux_weights(dit, pretrained_path) + + encoder_cfg = dict(model_config.get("encoder", {}) or cfg_dict.get("encoder", {}) or {}) + dtype = torch.bfloat16 + t5_encoder = None + clip_encoder = None + autoencoder = None + if encoder_cfg.get("t5_encoder"): + t5_encoder = HFEmbedder( + str(encoder_cfg["t5_encoder"]), + max_length=int(encoder_cfg.get("max_t5_length", 256)), + torch_dtype=dtype, + ) + if encoder_cfg.get("clip_encoder"): + clip_encoder = HFEmbedder( + str(encoder_cfg["clip_encoder"]), + max_length=int(encoder_cfg.get("max_clip_length", 77)), + torch_dtype=dtype, + ) + if encoder_cfg.get("autoencoder"): + ae_params = AutoEncoderParams( + resolution=int(encoder_cfg.get("resolution", 256)), + scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + autoencoder = load_autoencoder( + str(encoder_cfg["autoencoder"]), + ae_params, + dtype=dtype, + sample_z=bool(encoder_cfg.get("sample_z", True)), + ) + + training_cfg = FluxTrainingConfig( + model_variant=variant, + trainable_modules=cfg_dict.get("trainable_modules", "dit"), + guidance=float(cfg_dict.get("guidance", 1.0)), + autoencoder_scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + autoencoder_shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + pipeline = FluxFlowMatchTrainPipeline( + FluxFlowMatchTrainPipelineConfig( + autoencoder_scale_factor=training_cfg.autoencoder_scale_factor, + autoencoder_shift_factor=training_cfg.autoencoder_shift_factor, + guidance=training_cfg.guidance, + ) + ) + model = FluxForTraining( + dit=dit, + train_pipeline=pipeline, + model_config=training_cfg, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + raw_config={ + "model_config": model_config, + "flux_params": params.to_dict(), + }, + trainable_modules=training_cfg.trainable_modules, + ) + total_params, trainable_params = count_parameters(model) + logger.info(f"Built FLUX model: total={total_params:,} trainable={trainable_params:,}") + return model diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py index b0a728827..72bce2776 100644 --- a/primus/backends/diffusion/registry.py +++ b/primus/backends/diffusion/registry.py @@ -23,6 +23,18 @@ def _build_wan_dataset(dataset_config: dict): return build_wan_dataset(dataset_config) +def _build_flux_model(model_config: dict): + from primus.backends.diffusion.models.registrations.flux import build_flux_model + + return build_flux_model(model_config) + + +def _build_flux_dataset(dataset_config: dict): + from primus.backends.diffusion.data.registrations.flux import build_flux_dataset + + return build_flux_dataset(dataset_config) + + def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer @@ -35,9 +47,11 @@ def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = { + "flux": _build_flux_model, "wan": _build_wan_model, } DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = { + "flux": _build_flux_dataset, "wan": _build_wan_dataset, } TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = { diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py index ec5519574..cdd3f2f9b 100644 --- a/primus/backends/diffusion/schedulers/flow_match.py +++ b/primus/backends/diffusion/schedulers/flow_match.py @@ -59,6 +59,11 @@ def set_timesteps( if dynamic_shift_len is not None else self.exponential_shift_mu ) + if mu is None: + raise ValueError( + "`exponential_shift=True` requires either `dynamic_shift_len` " + "or `exponential_shift_mu`." + ) self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) else: self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py index d4df6355e..f1fbeeb37 100644 --- a/primus/backends/diffusion/trainers/__init__.py +++ b/primus/backends/diffusion/trainers/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Trainer registrations for the Primus Wan backend.""" +"""Trainer registrations for the Primus diffusion backend.""" from .fsdp2 import build_fsdp2_trainer diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py index c5d5227fd..89e3b764e 100644 --- a/primus/backends/diffusion/trainers/base.py +++ b/primus/backends/diffusion/trainers/base.py @@ -385,6 +385,9 @@ def _infer_batch_size_from_sequences(self, value) -> int | None: return None def _infer_local_batch_size(self, batch) -> int: + if isinstance(batch, (list, tuple)): + return len(batch) + tensor_batch_size = self._infer_batch_size_from_tensors(batch) if tensor_batch_size is not None: return tensor_batch_size @@ -470,6 +473,11 @@ def train(self): self.model.train() self.optimizer.zero_grad(set_to_none=True) torch.cuda.reset_peak_memory_stats() + debug_steps = os.getenv("DIFFUSION_DEBUG_STEPS", "0") == "1" + + def debug_phase(message: str) -> None: + if debug_steps: + logger.info(f"debug_step rank={self.rank} step={self.global_step} {message}") start_time = time.time() last_log_time = start_time @@ -485,21 +493,27 @@ def train(self): for batch_idx, batch in enumerate(self.dataloader): is_update_step = ((batch_idx + 1) % max(1, self.grad_accum_steps)) == 0 local_samples_in_update += self._infer_local_batch_size(batch) + debug_phase(f"batch_idx={batch_idx} before_compute_loss update={is_update_step}") with self._grad_sync_context(is_update_step): raw_loss = self.compute_loss(batch) + debug_phase("after_compute_loss") update_loss_sum += raw_loss.detach().float().item() update_loss_count += 1 loss = raw_loss / max(1, self.grad_accum_steps) loss.backward() + debug_phase("after_backward") if is_update_step: loss_val = update_loss_sum / max(1, update_loss_count) update_loss_sum = 0.0 update_loss_count = 0 + debug_phase("before_clip_grad_norm") grad_norm = self._clip_grad_norm() + debug_phase("before_optimizer_step") self.optimizer.step() + debug_phase("after_optimizer_step") self.lr_scheduler.step() self.optimizer.zero_grad(set_to_none=True) self.global_step += 1 diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py index 1fa79b177..e3bd72a70 100644 --- a/primus/backends/diffusion/trainers/fsdp2.py +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -370,6 +370,10 @@ def _save_dit(self, save_path: str) -> None: def save_model(self): """Save final model using the configured strategy.""" + if self.save_strategy in ("none", "skip", "disabled"): + if self.rank == 0: + logger.info(f"Skipping final checkpoint save (strategy={self.save_strategy})") + return if self.save_strategy == "dit_only": save_path = os.path.join(self.output_dir, "dit_model.safetensors") self._save_dit(save_path) diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py index 3d7aa5fe7..b4ecc5389 100644 --- a/primus/backends/diffusion/utils/__init__.py +++ b/primus/backends/diffusion/utils/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Utility package for the Primus Wan backend. +"""Utility package for the Primus diffusion backend. -Heavy vision helpers are imported lazily by the qwen_vl_utils video path. +Heavy vision helpers are imported lazily by the video dataset path. """ diff --git a/primus/configs/models/diffusion/flux.1_dev_t2i.yaml b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml new file mode 100644 index 000000000..64d5287bd --- /dev/null +++ b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml @@ -0,0 +1,18 @@ +model: + name: flux + config: + # Optional: set PRETRAINED_PATH to initialize from existing FLUX DiT weights. + load_from_pretrained_path: ${PRETRAINED_PATH:} + config: + model_variant: flux-dev + trainable_modules: dit + guidance: ${FLUX_GUIDANCE:1.0} + autoencoder_scale_factor: 0.3611 + autoencoder_shift_factor: 0.1159 + # Optional for precomputed data, required for raw image-text data. + encoder: + t5_encoder: ${T5_ENCODER:} + clip_encoder: ${CLIP_ENCODER:} + autoencoder: ${VAE_CHECKPOINT:} + max_t5_length: ${MAX_T5_LENGTH:256} + max_clip_length: 77 diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py index b23e6a5ae..92bf5e678 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -71,6 +71,42 @@ def _require_path(path: str | None, description: str, *, kind: str = "any") -> N _log(f"{description}: {resolved}") +def _require_optional_path(path: str | None, description: str, *, kind: str = "any") -> None: + if _is_placeholder(path): + _log(f"{description}: not configured; skipping optional initialization") + return + _require_path(path, description, kind=kind) + + +def _looks_like_local_path(value: str) -> bool: + if value.startswith(("/", "./", "../", "~")): + return True + parts = value.split("/") + return len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")) + + +def _looks_like_hf_reference(value: str) -> bool: + if _looks_like_local_path(value): + return False + parts = value.split("/") + return len(parts) >= 2 and all(part for part in parts[:2]) + + +def _validate_local_or_hf_id(value: str | None, description: str) -> None: + if _is_placeholder(value): + _fail(f"{description} is not configured: {value!r}") + value = str(value) + path = Path(value).expanduser() + if path.exists(): + _log(f"{description}: {path}") + elif _looks_like_local_path(value): + _fail(f"{description} path not found: {path}") + elif _looks_like_hf_reference(value): + _log(f"{description}: {value} (assuming Hugging Face reference)") + else: + _fail(f"{description} is neither an existing path nor a Hugging Face reference: {value!r}") + + def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: cfg = load_primus_config(config_path) selected_module = _select_module_name(cfg, module_name) @@ -87,16 +123,58 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) dataset_cfg = dataset.get("config", {}) processor_cfg = dataset_cfg.get("processor_config", {}) - encoder_cfg = model.get("config", {}).get("encoder", {}) or model.get("encoder", {}) - - _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") - _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") - _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") - model_cfg = model.get("config", {}) - _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") - _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") - _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + model_name = model.get("name") + + if model_name == "wan": + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") + _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") + _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") + _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") + _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") + _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + elif model_name == "flux": + dataset_type = str(dataset_cfg.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + dataset_name = dataset_cfg.get("dataset") + dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower() + if dataset_name == "cc12m-test" and _is_placeholder(dataset_cfg.get("dataset_path")): + _fail( + "FLUX raw dataset `cc12m-test` requires `dataset_path` to point at local webdataset tars" + ) + elif dataset_name == "cc12m-wds": + if _is_placeholder(dataset_cfg.get("dataset_path")): + _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") + else: + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + elif dataset_format == "hf_repo": + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + else: + kind = "file" if dataset_format == "jsonl" else "dir" + _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind) + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _validate_local_or_hf_id(encoder_cfg.get("t5_encoder"), "FLUX T5 encoder") + _validate_local_or_hf_id(encoder_cfg.get("clip_encoder"), "FLUX CLIP encoder") + _validate_local_or_hf_id(encoder_cfg.get("autoencoder"), "FLUX autoencoder checkpoint") + elif dataset_type == "precomputed": + _require_path(dataset_cfg.get("dataset_path"), "FLUX precomputed dataset", kind="dir") + prompt_dropout_prob = float(processor_cfg.get("prompt_dropout_prob", 0.0) or 0.0) + if prompt_dropout_prob > 0.0: + empty_dir = processor_cfg.get("empty_encodings_path") + _require_path(empty_dir, "FLUX empty encodings directory", kind="dir") + _require_path(str(Path(empty_dir) / "t5_empty.npy"), "FLUX empty T5 encoding", kind="file") + _require_path( + str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file" + ) + else: + _fail(f"unsupported FLUX dataset_type: {dataset_type!r}") + _require_optional_path( + model_cfg.get("load_from_pretrained_path"), + "FLUX DiT initialization checkpoint", + ) + else: + _fail(f"unsupported diffusion model for prepare hook: {model_name!r}") _log( "validated " diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index 38fabbeef..1e4547a8d 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -1,6 +1,9 @@ einops +datasets +huggingface_hub loguru safetensors +sentencepiece numpy pillow packaging diff --git a/tests/unit_tests/backends/__init__.py b/tests/unit_tests/backends/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/diffusion/__init__.py b/tests/unit_tests/backends/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py new file mode 100644 index 000000000..9a97445dd --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import pytest + +from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler + + +def test_exponential_shift_requires_mu_or_dynamic_shift_len(): + with pytest.raises(ValueError, match="exponential_shift=True"): + FlowMatchScheduler(exponential_shift=True) diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py new file mode 100644 index 000000000..b73e146c1 --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder +from primus.backends.diffusion.attention import ( + get_attention_backend, + set_attention_backend, +) +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.models.flux.math import apply_rope +from primus.backends.diffusion.models.flux.math import attention as flux_attention +from primus.backends.diffusion.models.flux.math import rope +from primus.backends.diffusion.models.registrations.flux import build_flux_model + + +def _finalize(params: dict): + builder = DiffusionArgBuilder() + builder.update(params) + return builder.finalize() + + +def test_flux_argument_builder_selects_flux_defaults(): + args = _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "training": {"steps": 7, "local_batch_size": 3}, + "data": { + "dataset_path": "/tmp/precomputed", + "dataset_type": "precomputed", + "empty_encodings_path": "/tmp/empty", + "prompt_dropout_prob": 0.25, + }, + } + ) + + assert args.model["name"] == "flux" + assert args.dataset["name"] == "flux" + assert args.dataset["config"]["dataset_path"] == "/tmp/precomputed" + assert args.dataset["config"]["processor_config"]["empty_encodings_path"] == "/tmp/empty" + assert args.dataset["config"]["processor_config"]["prompt_dropout_prob"] == 0.25 + assert args.trainer["args"]["max_steps"] == 7 + assert args.trainer["args"]["per_device_train_batch_size"] == 3 + assert args.trainer["args"]["attention_backend"] == "flash_attn_aiter" + assert args.trainer["args"]["fsdp_transformer_layer_cls_to_wrap"] == "DoubleStreamBlock,SingleStreamBlock" + + +def test_flux_argument_builder_maps_raw_dataset_type(): + args = _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "data": { + "dataset_type": "raw", + "dataset_format": "webdataset", + "dataset_path": "/tmp/cc12m_test", + "prompt_dropout_prob": 0.1, + "img_size": 128, + }, + } + ) + + assert args.dataset["config"]["dataset_type"] == "raw" + assert args.dataset["config"]["dataset_format"] == "webdataset" + assert args.dataset["config"]["processor_config"]["img_size"] == 128 + + +def test_flux_raw_dataset_name_defaults(): + with pytest.raises(ValueError, match="cc12m-test.*dataset_path"): + FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset") + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", "/tmp/cc12m_test", "webdataset") + assert path == "/tmp/cc12m_test" + assert fmt == "webdataset" + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-wds", None, "webdataset") + assert path == "pixparse/cc12m-wds" + assert fmt == "hf_repo" + + +def test_flux_argument_builder_rejects_sequence_parallelism(): + with pytest.raises(ValueError, match="sp_size"): + _finalize( + { + "model": {"name": "flux", "config": {"config": {"model_variant": "flux-dev"}}}, + "parallelism": {"sp_size": 2}, + } + ) + + +def test_flux_attention_dispatch_matches_sdpa_layout(): + previous_backend = get_attention_backend() + set_attention_backend("sdpa") + try: + torch.manual_seed(7) + q = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + k = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + v = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + pos = torch.arange(5, dtype=torch.float32).repeat(2, 1) + pe = rope(pos, dim=4, theta=10000).unsqueeze(1) + + actual = flux_attention(q, k, v, pe=pe) + q_rope, k_rope = apply_rope(q, k, pe) + expected = torch.nn.functional.scaled_dot_product_attention(q_rope, k_rope, v) + expected = expected.transpose(1, 2).flatten(2) + + torch.testing.assert_close(actual, expected) + finally: + set_attention_backend(previous_backend) + + +def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 3, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + ] + + out = processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + assert out["t5_encodings"].shape == (2, 3, 8) + assert out["clip_encodings"].shape == (2, 4) + assert torch.count_nonzero(out["t5_encodings"]) == 0 + assert torch.count_nonzero(out["clip_encodings"]) == 0 + + +def test_flux_precomputed_processor_rejects_mismatched_empty_encoding(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + # Empty T5 encoding has sequence length 5, but the batch samples use length 3. + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 5, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + } + ] + + with pytest.raises(ValueError, match="empty T5 encoding shape"): + processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + +def test_flux_raw_processor_prepares_images_and_prompts(): + from PIL import Image + + processor = FluxRawImageTextProcessor( + {"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False} + ) + image = Image.fromarray(np.full((6, 10, 3), 127, dtype=np.uint8)) + + out = processor.prepare_batch( + batch=[{"image": image, "prompt": "a test image"}], + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert out["image"].shape == (1, 3, 8, 8) + assert out["prompts"] == [""] + + +def test_tiny_flux_model_computes_precomputed_loss(): + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + batch = { + "t5_encodings": torch.randn(2, 3, 8), + "clip_encodings": torch.randn(2, 4), + "mean": torch.randn(2, 1, 2, 2), + "logvar": torch.zeros(2, 1, 2, 2), + } + + outputs = model.forward_train(batch) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) + + +def test_flux_position_ids_are_float32_regardless_of_model_dtype(): + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + # Force a low-precision compute dtype; position ids must stay float32 so RoPE + # grid indices are not corrupted (bf16 only represents integers up to 256). + model.dit = model.dit.to(dtype=torch.bfloat16) + + captured = {} + original_forward = model.dit.forward + + def capturing_forward(*args, **kwargs): + captured["img_ids_dtype"] = kwargs["img_ids"].dtype + captured["txt_ids_dtype"] = kwargs["txt_ids"].dtype + return original_forward(*args, **kwargs) + + model.dit.forward = capturing_forward + model.forward_train( + { + "t5_encodings": torch.randn(2, 3, 8, dtype=torch.bfloat16), + "clip_encodings": torch.randn(2, 4, dtype=torch.bfloat16), + "mean": torch.randn(2, 1, 2, 2, dtype=torch.bfloat16), + "logvar": torch.zeros(2, 1, 2, 2, dtype=torch.bfloat16), + } + ) + + assert captured["img_ids_dtype"] == torch.float32 + assert captured["txt_ids_dtype"] == torch.float32 + + +def test_tiny_flux_model_computes_raw_loss_with_dummy_encoders(): + class DummyAutoencoder(torch.nn.Module): + def encode(self, image): + return image[:, :1, :2, :2] + + class DummyT5(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 3, 8) + + class DummyClip(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 4) + + model = build_flux_model( + { + "config": { + "model_variant": "flux-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + model.autoencoder = DummyAutoencoder() + model.t5_encoder = DummyT5() + model.clip_encoder = DummyClip() + + outputs = model.forward_train( + { + "image": torch.randn(2, 3, 4, 4), + "prompts": ["cat", "dog"], + } + ) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) diff --git a/tests/unit_tests/backends/megatron/__init__.py b/tests/unit_tests/backends/megatron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/torchtitan/__init__.py b/tests/unit_tests/backends/torchtitan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/cli/__init__.py b/tests/unit_tests/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/__init__.py b/tests/unit_tests/core/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/backend/__init__.py b/tests/unit_tests/core/backend/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/launcher/__init__.py b/tests/unit_tests/core/launcher/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/patches/__init__.py b/tests/unit_tests/core/patches/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/runtime/__init__.py b/tests/unit_tests/core/runtime/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/trainer/__init__.py b/tests/unit_tests/core/trainer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/utils/__init__.py b/tests/unit_tests/core/utils/__init__.py new file mode 100644 index 000000000..e69de29bb