Add FastWAM policy

tests/policies/fastwam/test_fastwam_wan_components.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+import pytest
+from torch import nn
+
+from lerobot.policies.fastwam import modeling_fastwam
+from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+from lerobot.policies.fastwam.modeling_fastwam import (
+    FastWAMPolicy,
+    resolve_wan_component_paths,
+)
+from lerobot.policies.fastwam.wan_components import (
+    WAN_DIT_PATTERN,
+    WAN_T5_CHECKPOINT,
+    WAN_T5_TOKENIZER,
+    WAN_VAE_CHECKPOINT,
+    resolve_wan_checkpoint_paths,
+)
+
+
+def _make_wan_component_tree(root: Path) -> None:
+    tokenizer = root / WAN_T5_TOKENIZER
+    tokenizer.mkdir(parents=True)
+    (root / WAN_VAE_CHECKPOINT).touch()
+    (root / WAN_T5_CHECKPOINT).touch()
+    (root / "diffusion_pytorch_model-00001-of-00001.safetensors").touch()
+    (tokenizer / "tokenizer.json").touch()
+
+
+def test_resolve_wan_component_paths_finds_complete_local_directory(tmp_path):
+    _make_wan_component_tree(tmp_path)
+
+    paths = resolve_wan_component_paths(tmp_path)
+
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+
+
+def test_resolve_wan_component_paths_does_not_require_original_dit_shards(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    for shard in tmp_path.glob(WAN_DIT_PATTERN):
+        shard.unlink()
+
+    paths = resolve_wan_component_paths(tmp_path)
+
+    assert paths.dit == []
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+
+
+def test_resolve_wan_checkpoint_paths_uses_official_wan_layout(tmp_path):
+    _make_wan_component_tree(tmp_path)
+
+    paths = resolve_wan_checkpoint_paths(tmp_path)
+
+    assert paths.root == tmp_path
+    assert paths.dit == [tmp_path / "diffusion_pytorch_model-00001-of-00001.safetensors"]
+    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
+    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
+    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
+    assert WAN_DIT_PATTERN == "diffusion_pytorch_model*.safetensors"
+
+
+def test_resolve_wan_component_paths_rejects_partial_local_directory(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    (tmp_path / WAN_T5_CHECKPOINT).unlink()
+
+    with pytest.raises(FileNotFoundError, match="text encoder"):
+        resolve_wan_component_paths(tmp_path)
+
+
+def test_policy_config_construction_loads_wan22_backbone_from_config(monkeypatch):
+    class TinyCore(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.text_encoder = None
+
+    calls = []
+
+    def fake_from_wan22_pretrained(**kwargs):
+        calls.append(kwargs)
+        return TinyCore()
+
+    monkeypatch.setattr(
+        "lerobot.policies.fastwam.modular_fastwam.FastWAM.from_wan22_pretrained",
+        fake_from_wan22_pretrained,
+    )
+
+    cfg = FastWAMConfig()
+    policy = FastWAMPolicy(cfg)
+
+    assert policy.model.text_encoder is None
+    assert calls == [
+        {
+            "device": cfg.device,
+            "torch_dtype": modeling_fastwam._dtype_from_name(cfg.torch_dtype),
+            "model_id": "Wan-AI/Wan2.2-TI2V-5B",
+            "tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B",
+            "tokenizer_max_len": cfg.tokenizer_max_len,
+            "load_text_encoder": cfg.load_text_encoder,
+            "proprio_dim": cfg.proprio_dim,
+            "video_dit_config": cfg.video_dit_config,
+            "action_dit_config": cfg.action_dit_config,
+            "mot_checkpoint_mixed_attn": cfg.mot_checkpoint_mixed_attn,
+            "video_train_shift": float(cfg.video_scheduler["train_shift"]),
+            "video_infer_shift": float(cfg.video_scheduler["infer_shift"]),
+            "video_num_train_timesteps": int(cfg.video_scheduler["num_train_timesteps"]),
+            "action_train_shift": float(cfg.action_scheduler["train_shift"]),
+            "action_infer_shift": float(cfg.action_scheduler["infer_shift"]),
+            "action_num_train_timesteps": int(cfg.action_scheduler["num_train_timesteps"]),
+            "loss_lambda_video": float(cfg.loss["lambda_video"]),
+            "loss_lambda_action": float(cfg.loss["lambda_action"]),
+        }
+    ]
+
+
+def test_explicit_local_wan_path_is_preserved(tmp_path):
+    cfg = FastWAMConfig(model_id=str(tmp_path), tokenizer_model_id=str(tmp_path))
+
+    assert cfg.model_id == str(tmp_path)
+    assert cfg.tokenizer_model_id == str(tmp_path)
+
+
+def test_other_hub_model_ids_are_rejected():
+    with pytest.raises(ValueError, match="model_id"):
+        FastWAMConfig(model_id="somebody/other-model")
+
+    with pytest.raises(ValueError, match="tokenizer_model_id"):
+        FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer")
+
+
+def test_resolve_wan_checkpoint_paths_can_skip_text_encoder(tmp_path):
+    _make_wan_component_tree(tmp_path)
+    (tmp_path / WAN_T5_CHECKPOINT).unlink()
+    shutil_tokenizer = tmp_path / WAN_T5_TOKENIZER
+    for child in shutil_tokenizer.iterdir():
+        child.unlink()
+    shutil_tokenizer.rmdir()
+    shutil_tokenizer.parent.rmdir()
+
+    paths = resolve_wan_checkpoint_paths(tmp_path, load_text_encoder=False)
+
+    assert paths.text_encoder is None
+    assert paths.tokenizer is None