diff --git a/README.md b/README.md index febbfc386..51bdc9974 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ We believe that a well-developed open-source code framework can lower the thresh ## Update History +- **August 18, 2026** We add optional Echo-Memory overlay support for Wan 2.1 1.3B, including standard inference, low VRAM inference, full training, and LoRA training scripts. For details, please refer to the [documentation](/docs/en/Model_Details/Wan.md) and [example code](/examples/wanvideo/). > DiffSynth-Studio has undergone major version updates, and some old features are no longer maintained. If you need to use old features, please switch to the [last historical version](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3) before the major version update. > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand. @@ -1372,6 +1373,7 @@ graph LR; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-I2V-14B-480P; Wan-AI/Wan2.1-I2V-14B-480P-->Wan-AI/Wan2.1-I2V-14B-720P; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-FLF2V-14B-720P; + Wan-AI/Wan2.1-T2V-1.3B-->Echo-Team/Echo-Memory; Wan-AI/Wan2.1-T2V-1.3B-->iic/VACE-Wan2.1-1.3B-Preview; iic/VACE-Wan2.1-1.3B-Preview-->Wan-AI/Wan2.1-VACE-1.3B; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-VACE-14B; @@ -1416,6 +1418,7 @@ Example code for Wan is available at: [/examples/wanvideo/](/examples/wanvideo/) | Model ID | Extra Inputs | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training | |-|-|-|-|-|-|-|-| |[Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)||[code](/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py)| +|[Echo-Team/Echo-Memory: context_k1](https://huggingface.co/Echo-Team/Echo-Memory)||[code](/examples/wanvideo/model_inference/Echo-Memory.py)|[code](/examples/wanvideo/model_inference_low_vram/Echo-Memory.py)|[code](/examples/wanvideo/model_training/full/Echo-Memory.sh)|[code](/examples/wanvideo/model_training/validate_full/Echo-Memory.py)|[code](/examples/wanvideo/model_training/lora/Echo-Memory.sh)|[code](/examples/wanvideo/model_training/validate_lora/Echo-Memory.py)| |[Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)||[code](/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py)| |[Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)|`input_image`|[code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py)| |[Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)|`input_image`|[code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py)| diff --git a/README_zh.md b/README_zh.md index c31c891f3..57d56b12a 100644 --- a/README_zh.md +++ b/README_zh.md @@ -32,6 +32,7 @@ DiffSynth 目前包括两个开源项目： ## 更新历史 +- **2026年8月18日** 我们为 Wan 2.1 1.3B 新增了可选 Echo-Memory 权重叠加支持，包含标准推理、低显存推理、全量训练和 LoRA 训练脚本。详情请参考[文档](/docs/zh/Model_Details/Wan.md)和[示例代码](/examples/wanvideo/)。 > DiffSynth-Studio 经历了大版本更新，部分旧功能已停止维护，如需使用旧版功能，请切换到大版本更新前的[最后一个历史版本](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3)。 > 目前本项目的开发人员有限，大部分工作由 [Artiprocher](https://github.com/Artiprocher) 和 [mi804](https://github.com/mi804) 负责，因此新功能的开发进展会比较缓慢，issue 的回复和解决速度有限，我们对此感到非常抱歉，请各位开发者理解。 @@ -1372,6 +1373,7 @@ graph LR; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-I2V-14B-480P; Wan-AI/Wan2.1-I2V-14B-480P-->Wan-AI/Wan2.1-I2V-14B-720P; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-FLF2V-14B-720P; + Wan-AI/Wan2.1-T2V-1.3B-->Echo-Team/Echo-Memory; Wan-AI/Wan2.1-T2V-1.3B-->iic/VACE-Wan2.1-1.3B-Preview; iic/VACE-Wan2.1-1.3B-Preview-->Wan-AI/Wan2.1-VACE-1.3B; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-VACE-14B; @@ -1416,6 +1418,7 @@ Wan 的示例代码位于：[/examples/wanvideo/](/examples/wanvideo/) |模型 ID|额外参数|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| |-|-|-|-|-|-|-|-| |[Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)||[code](/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py)| +|[Echo-Team/Echo-Memory: context_k1](https://huggingface.co/Echo-Team/Echo-Memory)||[code](/examples/wanvideo/model_inference/Echo-Memory.py)|[code](/examples/wanvideo/model_inference_low_vram/Echo-Memory.py)|[code](/examples/wanvideo/model_training/full/Echo-Memory.sh)|[code](/examples/wanvideo/model_training/validate_full/Echo-Memory.py)|[code](/examples/wanvideo/model_training/lora/Echo-Memory.sh)|[code](/examples/wanvideo/model_training/validate_lora/Echo-Memory.py)| |[Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)||[code](/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py)| |[Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)|`input_image`|[code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py)| |[Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)|`input_image`|[code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh)|[code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py)|[code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh)|[code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py)| diff --git a/diffsynth/pipelines/wan_video_echo_memory.py b/diffsynth/pipelines/wan_video_echo_memory.py new file mode 100644 index 000000000..7791d5838 --- /dev/null +++ b/diffsynth/pipelines/wan_video_echo_memory.py @@ -0,0 +1,115 @@ +from typing import Iterable, Optional + +from ..core import ModelConfig, load_state_dict + + +DEFAULT_REPO_ID = "Echo-Team/Echo-Memory" +DEFAULT_FILENAME = "context_k1/epoch-0.safetensors" +DEFAULT_DOWNLOAD_SOURCE = "huggingface" +SKIP_SUBSTRINGS = ( + "action_mlp", + "self_attn_with_action", + "block_wise_ssm", + "videossm_hybrid", + "spatial_memory_module", +) + + +def filter_dit_state_dict(state_dict: dict, skip_substrings: Iterable[str] = SKIP_SUBSTRINGS) -> dict: + skip_substrings = tuple(skip_substrings) + return { + key: value + for key, value in state_dict.items() + if not any(token in key for token in skip_substrings) + } + + +def _resolve_model_path( + model_config: Optional[ModelConfig] = None, + local_path: Optional[str] = None, + repo_id: str = DEFAULT_REPO_ID, + filename: str = DEFAULT_FILENAME, + download_source: str = DEFAULT_DOWNLOAD_SOURCE, +): + if local_path is not None: + return local_path + if model_config is None: + model_config = ModelConfig(model_id=repo_id, origin_file_pattern=filename, download_source=download_source) + model_config.download_if_necessary() + if isinstance(model_config.path, list): + if len(model_config.path) != 1: + raise ValueError(f"Expected exactly one Echo-Memory checkpoint, got {len(model_config.path)} files.") + return model_config.path[0] + return model_config.path + + +def _canonical_key(key: str) -> str: + return ".".join(part for part in key.split(".") if part != "module") + + +def _source_key_candidates(key: str): + yield key + for prefix in ("pipe.dit.", "dit."): + if key.startswith(prefix): + yield key[len(prefix):] + + +def _has_meta_tensor(state_dict: dict) -> bool: + return any(getattr(value, "is_meta", False) for value in state_dict.values()) + + +def align_dit_state_dict_to_model(state_dict: dict, model_state_dict: dict) -> tuple[dict, list[str]]: + target_key_map = {} + for key in model_state_dict: + target_key_map.setdefault(_canonical_key(key), []).append(key) + aligned, unaligned = {}, [] + for key, value in state_dict.items(): + matches = [] + for candidate in _source_key_candidates(key): + for target_key in target_key_map.get(_canonical_key(candidate), []): + if getattr(model_state_dict[target_key], "shape", None) == getattr(value, "shape", None): + matches.append(target_key) + matches = sorted(set(matches)) + if len(matches) == 1: + aligned[matches[0]] = value + else: + unaligned.append(key) + return aligned, unaligned + + +def load_echo_memory_dit( + pipe, + model_config: Optional[ModelConfig] = None, + repo_id: str = DEFAULT_REPO_ID, + filename: str = DEFAULT_FILENAME, + local_path: Optional[str] = None, + download_source: str = DEFAULT_DOWNLOAD_SOURCE, + torch_dtype=None, +): + if getattr(pipe, "dit", None) is None: + raise ValueError("pipe.dit is empty; load Wan 2.1 1.3B before overlaying Echo-Memory.") + + ckpt_path = _resolve_model_path( + model_config=model_config, + local_path=local_path, + repo_id=repo_id, + filename=filename, + download_source=download_source, + ) + raw = load_state_dict(ckpt_path, torch_dtype=torch_dtype, device="cpu") + filtered = filter_dit_state_dict(raw) + model_state_dict = pipe.dit.state_dict() + aligned, unaligned = align_dit_state_dict_to_model(filtered, model_state_dict) + missing, unexpected = pipe.dit.load_state_dict(aligned, strict=False, assign=_has_meta_tensor(model_state_dict)) + print( + f"[Echo-Memory] overlaid {len(aligned)}/{len(raw)} DiT keys from {ckpt_path} " + f"(skipped={len(raw) - len(filtered)}, unaligned={len(unaligned)}, " + f"missing={len(missing)}, unexpected={len(unexpected)})" + ) + if unaligned: + print("[Echo-Memory] unaligned example:", ", ".join(sorted(unaligned)[:5])) + if missing: + print("[Echo-Memory] missing example:", ", ".join(sorted(missing)[:5])) + if unexpected: + print("[Echo-Memory] unexpected example:", ", ".join(sorted(unexpected)[:5])) + return missing, unexpected diff --git a/docs/en/Model_Details/Wan.md b/docs/en/Model_Details/Wan.md index 632f007b8..28a2d8041 100644 --- a/docs/en/Model_Details/Wan.md +++ b/docs/en/Model_Details/Wan.md @@ -70,6 +70,7 @@ graph LR; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-I2V-14B-480P; Wan-AI/Wan2.1-I2V-14B-480P-->Wan-AI/Wan2.1-I2V-14B-720P; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-FLF2V-14B-720P; + Wan-AI/Wan2.1-T2V-1.3B-->Echo-Team/Echo-Memory; Wan-AI/Wan2.1-T2V-1.3B-->iic/VACE-Wan2.1-1.3B-Preview; iic/VACE-Wan2.1-1.3B-Preview-->Wan-AI/Wan2.1-VACE-1.3B; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-VACE-14B; @@ -108,6 +109,7 @@ graph LR; | Model ID | Extra Inputs | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training | |-|-|-|-|-|-|-|-| |[Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py)| +|[Echo-Team/Echo-Memory: context_k1](https://huggingface.co/Echo-Team/Echo-Memory)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Echo-Memory.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Echo-Memory.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Echo-Memory.py)| |[Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py)| |[Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)|`input_image`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py)| |[Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)|`input_image`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py)| @@ -145,6 +147,18 @@ graph LR; |[Wan-AI/Wan-Dancer-14B (global model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B)|`wantodance_music_path`, `wantodance_reference_image`, `wantodance_fps`, `wantodance_keyframes`, `wantodance_keyframes_mask`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan-Dancer-14B-global.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan-Dancer-14B-global.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-global.py)| |[Wan-AI/Wan-Dancer-14B (local model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B)|`wantodance_music_path`, `wantodance_reference_image`, `wantodance_fps`, `wantodance_keyframes`, `wantodance_keyframes_mask`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan-Dancer-14B-local.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan-Dancer-14B-local.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-local.py)| +### Community: Echo-Memory + +[Echo-Memory](https://github.com/Echo-Team-Joy-Future-Academy-JD/Echo-Memory) ([arXiv:2606.09803](https://arxiv.org/abs/2606.09803)) is a controlled memory study on Wan 2.1 1.3B. The released `context_k1` checkpoint can be overlaid on `Wan-AI/Wan2.1-T2V-1.3B` through the optional loader in [`diffsynth.pipelines.wan_video_echo_memory`](https://github.com/modelscope/DiffSynth-Studio/blob/main/diffsynth/pipelines/wan_video_echo_memory.py). + +* Weights: [Echo-Team/Echo-Memory](https://huggingface.co/Echo-Team/Echo-Memory) +* Inference: [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Echo-Memory.py) +* Low VRAM inference: [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Echo-Memory.py) +* Full training: [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Echo-Memory.sh) +* LoRA training: [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Echo-Memory.sh) + +The current integration overlays matching DiT weights onto the Wan DiT. Camera-action MLP, block-wise SSM, spatial memory, and the multi-chunk revisit protocol remain in the upstream Echo-Memory repository. + * FP8 Precision Training: [doc](../Training/FP8_Precision.md), [code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/fp8_training/) * Two-stage Split Training: [doc](../Training/Split_Training.md), [code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/split_training/) * End-to-end Direct Distillation: [doc](../Training/Direct_Distill.md), [code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/direct_distill/) @@ -304,6 +318,9 @@ Wan series models are uniformly trained through [`examples/wanvideo/model_traini * Wan Series Specific Parameters * `--tokenizer_path`: Path of the tokenizer, applicable to text-to-video models, leave blank to automatically download from remote. * `--audio_processor_path`: Path of the audio processor, applicable to speech-to-video models, leave blank to automatically download from remote. + * `--echo_memory_path`: Local path to the Echo-Memory DiT checkpoint. + * `--echo_memory_model_id_with_origin_path`: Echo-Memory model ID with origin file path, e.g., `"Echo-Team/Echo-Memory:context_k1/epoch-0.safetensors"`. + * `--echo_memory_download_source`: Download source for the Echo-Memory checkpoint, default is `"huggingface"`. We have built a sample video dataset for your testing. You can download this dataset with the following command: diff --git a/docs/zh/Model_Details/Wan.md b/docs/zh/Model_Details/Wan.md index 55fd0885f..7d0959ae6 100644 --- a/docs/zh/Model_Details/Wan.md +++ b/docs/zh/Model_Details/Wan.md @@ -70,6 +70,7 @@ graph LR; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-I2V-14B-480P; Wan-AI/Wan2.1-I2V-14B-480P-->Wan-AI/Wan2.1-I2V-14B-720P; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-FLF2V-14B-720P; + Wan-AI/Wan2.1-T2V-1.3B-->Echo-Team/Echo-Memory; Wan-AI/Wan2.1-T2V-1.3B-->iic/VACE-Wan2.1-1.3B-Preview; iic/VACE-Wan2.1-1.3B-Preview-->Wan-AI/Wan2.1-VACE-1.3B; Wan-AI/Wan2.1-T2V-14B-->Wan-AI/Wan2.1-VACE-14B; @@ -109,6 +110,7 @@ graph LR; |模型 ID|额外参数|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| |-|-|-|-|-|-|-|-| |[Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py)| +|[Echo-Team/Echo-Memory: context_k1](https://huggingface.co/Echo-Team/Echo-Memory)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Echo-Memory.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Echo-Memory.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Echo-Memory.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Echo-Memory.py)| |[Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py)| |[Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)|`input_image`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py)| |[Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)|`input_image`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py)| @@ -146,6 +148,18 @@ graph LR; |[Wan-AI/Wan-Dancer-14B (global model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B)|`wantodance_music_path`, `wantodance_reference_image`, `wantodance_fps`, `wantodance_keyframes`, `wantodance_keyframes_mask`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan-Dancer-14B-global.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-global.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan-Dancer-14B-global.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-global.py)| |[Wan-AI/Wan-Dancer-14B (local model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B)|`wantodance_music_path`, `wantodance_reference_image`, `wantodance_fps`, `wantodance_keyframes`, `wantodance_keyframes_mask`|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Wan-Dancer-14B-local.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-local.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Wan-Dancer-14B-local.sh)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-local.py)| +### 社区：Echo-Memory + +[Echo-Memory](https://github.com/Echo-Team-Joy-Future-Academy-JD/Echo-Memory)（[arXiv:2606.09803](https://arxiv.org/abs/2606.09803)）是在 Wan 2.1 1.3B 上的受控 memory 研究。公开的 `context_k1` 权重可通过 [`diffsynth.pipelines.wan_video_echo_memory`](https://github.com/modelscope/DiffSynth-Studio/blob/main/diffsynth/pipelines/wan_video_echo_memory.py) 中的可选加载器叠加到 `Wan-AI/Wan2.1-T2V-1.3B`。 + +* 权重：[Echo-Team/Echo-Memory](https://huggingface.co/Echo-Team/Echo-Memory) +* 推理：[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference/Echo-Memory.py) +* 低显存推理：[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_inference_low_vram/Echo-Memory.py) +* 全量训练：[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/full/Echo-Memory.sh) +* LoRA 训练：[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/wanvideo/model_training/lora/Echo-Memory.sh) + +当前接入会将匹配的 DiT 权重覆盖到 Wan DiT；相机动作 MLP、block-wise SSM、spatial memory 和多 chunk revisit 协议仍在 Echo-Memory 仓库中。 + * FP8 精度训练：[doc](../Training/FP8_Precision.md)、[code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/fp8_training/) * 两阶段拆分训练：[doc](../Training/Split_Training.md)、[code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/split_training/) * 端到端直接蒸馏：[doc](../Training/Direct_Distill.md)、[code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo/model_training/special/direct_distill/) @@ -305,6 +319,9 @@ Wan 系列模型统一通过 [`examples/wanvideo/model_training/train.py`](https * Wan 系列专有参数 * `--tokenizer_path`: tokenizer 的路径，适用于文生视频模型，留空则自动从远程下载。 * `--audio_processor_path`: 音频处理器的路径，适用于语音到视频模型，留空则自动从远程下载。 + * `--echo_memory_path`: Echo-Memory DiT 权重的本地路径。 + * `--echo_memory_model_id_with_origin_path`: Echo-Memory 权重的模型 ID 与原始文件路径，例如 `"Echo-Team/Echo-Memory:context_k1/epoch-0.safetensors"`。 + * `--echo_memory_download_source`: Echo-Memory 权重的下载源，默认值为 `"huggingface"`。 我们构建了一个样例视频数据集，以方便您进行测试，通过以下命令可以下载这个数据集： diff --git a/examples/wanvideo/model_inference/Echo-Memory.py b/examples/wanvideo/model_inference/Echo-Memory.py new file mode 100644 index 000000000..593a91345 --- /dev/null +++ b/examples/wanvideo/model_inference/Echo-Memory.py @@ -0,0 +1,25 @@ +import torch +from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig +from diffsynth.pipelines.wan_video_echo_memory import load_echo_memory_dit +from diffsynth.utils.data import save_video + + +pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth"), + ], + tokenizer_config=ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/"), +) + +load_echo_memory_dit(pipe) + +video = pipe( + prompt="纪实摄影风格画面，一只活泼的小狗在绿茵茵的草地上迅速奔跑。小狗毛色棕黄，两只耳朵立起，神情专注而欢快。阳光洒在它身上，使得毛发看上去格外柔软而闪亮。背景是一片开阔的草地，偶尔点缀着几朵野花，远处隐约可见蓝天和几片白云。透视感鲜明，捕捉小狗奔跑时的动感和四周草地的生机。中景侧面移动视角。", + negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + seed=0, tiled=True, +) +save_video(video, "video_Echo-Memory-context_k1.mp4", fps=15, quality=5) diff --git a/examples/wanvideo/model_inference_low_vram/Echo-Memory.py b/examples/wanvideo/model_inference_low_vram/Echo-Memory.py new file mode 100644 index 000000000..699e8c0d4 --- /dev/null +++ b/examples/wanvideo/model_inference_low_vram/Echo-Memory.py @@ -0,0 +1,36 @@ +import torch +from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig +from diffsynth.pipelines.wan_video_echo_memory import load_echo_memory_dit +from diffsynth.utils.data import save_video + + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.bfloat16, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", **vram_config), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", **vram_config), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2, +) + +load_echo_memory_dit(pipe) + +video = pipe( + prompt="纪实摄影风格画面，一只活泼的小狗在绿茵茵的草地上迅速奔跑。小狗毛色棕黄，两只耳朵立起，神情专注而欢快。阳光洒在它身上，使得毛发看上去格外柔软而闪亮。背景是一片开阔的草地，偶尔点缀着几朵野花，远处隐约可见蓝天和几片白云。透视感鲜明，捕捉小狗奔跑时的动感和四周草地的生机。中景侧面移动视角。", + negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + seed=0, tiled=True, +) +save_video(video, "video_Echo-Memory-context_k1.mp4", fps=15, quality=5) diff --git a/examples/wanvideo/model_training/full/Echo-Memory.sh b/examples/wanvideo/model_training/full/Echo-Memory.sh new file mode 100755 index 000000000..33660dd53 --- /dev/null +++ b/examples/wanvideo/model_training/full/Echo-Memory.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "wanvideo/Wan2.1-T2V-1.3B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/wanvideo/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/wanvideo/Wan2.1-T2V-1.3B \ + --dataset_metadata_path data/diffsynth_example_dataset/wanvideo/Wan2.1-T2V-1.3B/metadata.csv \ + --data_file_keys "video" \ + --height 480 \ + --width 832 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.1-T2V-1.3B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.1-T2V-1.3B:Wan2.1_VAE.pth" \ + --echo_memory_model_id_with_origin_path "Echo-Team/Echo-Memory:context_k1/epoch-0.safetensors" \ + --echo_memory_download_source "huggingface" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/Echo-Memory_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/wanvideo/model_training/lora/Echo-Memory.sh b/examples/wanvideo/model_training/lora/Echo-Memory.sh new file mode 100755 index 000000000..57f3783b0 --- /dev/null +++ b/examples/wanvideo/model_training/lora/Echo-Memory.sh @@ -0,0 +1,20 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "wanvideo/Wan2.1-T2V-1.3B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/wanvideo/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/wanvideo/Wan2.1-T2V-1.3B \ + --dataset_metadata_path data/diffsynth_example_dataset/wanvideo/Wan2.1-T2V-1.3B/metadata.csv \ + --data_file_keys "video" \ + --height 480 \ + --width 832 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.1-T2V-1.3B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.1-T2V-1.3B:Wan2.1_VAE.pth" \ + --echo_memory_model_id_with_origin_path "Echo-Team/Echo-Memory:context_k1/epoch-0.safetensors" \ + --echo_memory_download_source "huggingface" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --output_path "./models/train/Echo-Memory_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "q,k,v,o,ffn.0,ffn.2" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/wanvideo/model_training/train.py b/examples/wanvideo/model_training/train.py index 45c941356..fc3467b99 100644 --- a/examples/wanvideo/model_training/train.py +++ b/examples/wanvideo/model_training/train.py @@ -2,6 +2,7 @@ from diffsynth.core import UnifiedDataset from diffsynth.core.data.operators import LoadVideo, LoadAudio, ImageCropAndResize, ToAbsolutePath from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig +from diffsynth.pipelines.wan_video_echo_memory import load_echo_memory_dit from diffsynth.diffusion import * os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -11,6 +12,7 @@ def __init__( self, model_paths=None, model_id_with_origin_paths=None, tokenizer_path=None, audio_processor_path=None, + echo_memory_path=None, echo_memory_model_id_with_origin_path=None, echo_memory_download_source="huggingface", trainable_models=None, lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None, preset_lora_path=None, preset_lora_model=None, @@ -36,6 +38,11 @@ def __init__( tokenizer_config = ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/") if tokenizer_path is None else ModelConfig(tokenizer_path) audio_processor_config = self.parse_path_or_model_id(audio_processor_path) self.pipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config, audio_processor_config=audio_processor_config) + if echo_memory_path is not None or echo_memory_model_id_with_origin_path is not None: + echo_memory_config = self.parse_path_or_model_id(echo_memory_model_id_with_origin_path) if echo_memory_model_id_with_origin_path is not None else None + if echo_memory_config is not None: + echo_memory_config.download_source = echo_memory_download_source + load_echo_memory_dit(self.pipe, model_config=echo_memory_config, local_path=echo_memory_path) self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model) self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt) @@ -119,6 +126,9 @@ def wan_parser(): parser = add_video_size_config(parser) parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.") parser.add_argument("--audio_processor_path", type=str, default=None, help="Path to the audio processor. If provided, the processor will be used for Wan2.2-S2V model.") + parser.add_argument("--echo_memory_path", type=str, default=None, help="Path to the Echo-Memory DiT checkpoint.") + parser.add_argument("--echo_memory_model_id_with_origin_path", type=str, default=None, help="Echo-Memory checkpoint in model_id:origin_file_pattern format.") + parser.add_argument("--echo_memory_download_source", type=str, default="huggingface", help="Download source for the Echo-Memory checkpoint.") parser.add_argument("--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).") parser.add_argument("--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).") parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") @@ -133,11 +143,19 @@ def wan_parser(): gradient_accumulation_steps=args.gradient_accumulation_steps, kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)], ) + data_file_keys = args.data_file_keys.split(",") + extra_inputs = args.extra_inputs.split(",") if args.extra_inputs is not None else [] + special_operator_map = { + "animate_face_video": ToAbsolutePath(args.dataset_base_path) >> LoadVideo(args.num_frames, 4, 1, frame_processor=ImageCropAndResize(512, 512, None, 16, 16)), + "wantodance_music_path": ToAbsolutePath(args.dataset_base_path), + } + if "input_audio" in data_file_keys or "input_audio" in extra_inputs: + special_operator_map["input_audio"] = ToAbsolutePath(args.dataset_base_path) >> LoadAudio(sr=16000) dataset = UnifiedDataset( base_path=args.dataset_base_path, metadata_path=args.dataset_metadata_path, repeat=args.dataset_repeat, - data_file_keys=args.data_file_keys.split(","), + data_file_keys=data_file_keys, main_data_operator=UnifiedDataset.default_video_operator( base_path=args.dataset_base_path, max_pixels=args.max_pixels, @@ -149,17 +167,16 @@ def wan_parser(): time_division_factor=4 if not args.framewise_decoding else 1, time_division_remainder=1 if not args.framewise_decoding else 0, ), - special_operator_map={ - "animate_face_video": ToAbsolutePath(args.dataset_base_path) >> LoadVideo(args.num_frames, 4, 1, frame_processor=ImageCropAndResize(512, 512, None, 16, 16)), - "input_audio": ToAbsolutePath(args.dataset_base_path) >> LoadAudio(sr=16000), - "wantodance_music_path": ToAbsolutePath(args.dataset_base_path), - } + special_operator_map=special_operator_map, ) model = WanTrainingModule( model_paths=args.model_paths, model_id_with_origin_paths=args.model_id_with_origin_paths, tokenizer_path=args.tokenizer_path, audio_processor_path=args.audio_processor_path, + echo_memory_path=args.echo_memory_path, + echo_memory_model_id_with_origin_path=args.echo_memory_model_id_with_origin_path, + echo_memory_download_source=args.echo_memory_download_source, trainable_models=args.trainable_models, lora_base_model=args.lora_base_model, lora_target_modules=args.lora_target_modules, diff --git a/examples/wanvideo/model_training/validate_full/Echo-Memory.py b/examples/wanvideo/model_training/validate_full/Echo-Memory.py new file mode 100644 index 000000000..b516ca84f --- /dev/null +++ b/examples/wanvideo/model_training/validate_full/Echo-Memory.py @@ -0,0 +1,26 @@ +import torch +from diffsynth.core import load_state_dict +from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig +from diffsynth.pipelines.wan_video_echo_memory import load_echo_memory_dit +from diffsynth.utils.data import save_video + +pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"), + ], +) +load_echo_memory_dit(pipe) +state_dict = load_state_dict("models/train/Echo-Memory_full/epoch-1.safetensors") +pipe.dit.load_state_dict(state_dict) +pipe.enable_vram_management(num_persistent_param_in_dit=None) +video = pipe( + prompt="纪实摄影风格画面，一只活泼的小狗在绿茵茵的草地上迅速奔跑。小狗毛色棕黄，两只耳朵立起，神情专注而欢快。阳光洒在它身上，使得毛发看上去格外柔软而闪亮。背景是一片开阔的草地，偶尔有几朵野花点缀，远处隐约可见蓝天和树木。小狗奔跑时四肢舒展，尾巴微微上扬，展现出轻快灵动的姿态。中景跟拍视角。", + negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + seed=0, + tiled=True, +) +save_video(video, "video_Echo-Memory_full.mp4", fps=15, quality=5) diff --git a/examples/wanvideo/model_training/validate_lora/Echo-Memory.py b/examples/wanvideo/model_training/validate_lora/Echo-Memory.py new file mode 100644 index 000000000..7cd6359e6 --- /dev/null +++ b/examples/wanvideo/model_training/validate_lora/Echo-Memory.py @@ -0,0 +1,24 @@ +import torch +from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig +from diffsynth.pipelines.wan_video_echo_memory import load_echo_memory_dit +from diffsynth.utils.data import save_video + +pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"), + ], +) +load_echo_memory_dit(pipe) +pipe.load_lora(pipe.dit, "models/train/Echo-Memory_lora/epoch-4.safetensors", alpha=1) +pipe.enable_vram_management(num_persistent_param_in_dit=None) +video = pipe( + prompt="纪实摄影风格画面，一只活泼的小狗在绿茵茵的草地上迅速奔跑。小狗毛色棕黄，两只耳朵立起，神情专注而欢快。阳光洒在它身上，使得毛发看上去格外柔软而闪亮。背景是一片开阔的草地，偶尔有几朵野花点缀，远处隐约可见蓝天和树木。小狗奔跑时四肢舒展，尾巴微微上扬，展现出轻快灵动的姿态。中景跟拍视角。", + negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + seed=0, + tiled=True, +) +save_video(video, "video_Echo-Memory_lora.mp4", fps=15, quality=5)