Polish Echo-Memory Wan video integration

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)|

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)|

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