mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-10 03:21:54 +00:00
Add FastWAM policy review updates
This commit is contained in:
ZibinDong
committed by
Maxime Ellerbach
Maxime Ellerbach
parent
a343ed3a63
commit
dfc0170b4d
@@ -1,254 +0,0 @@
|
||||
#!/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.
|
||||
|
||||
import inspect
|
||||
import typing
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from lerobot.configs import FeatureType, PolicyFeature
|
||||
from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
|
||||
from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
|
||||
from lerobot.policies.fastwam.processor_fastwam import make_fastwam_pre_post_processors
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
|
||||
def test_package_init_exports_required_symbols():
|
||||
init_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "__init__.py").read_text()
|
||||
|
||||
assert "FastWAMConfig" in init_source
|
||||
assert "make_fastwam_pre_post_processors" in init_source
|
||||
|
||||
|
||||
def test_policy_config_is_exported_from_public_policies_package():
|
||||
import lerobot.policies as policies
|
||||
|
||||
assert policies.FastWAMConfig is FastWAMConfig
|
||||
assert "FastWAMConfig" in policies.__all__
|
||||
|
||||
|
||||
def test_fastwam_policy_docs_are_registered():
|
||||
readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "README.md"
|
||||
wan_readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan" / "README.md"
|
||||
policy_readme_path = ROOT / "docs" / "source" / "policy_fastwam_README.md"
|
||||
guide_path = ROOT / "docs" / "source" / "fastwam.mdx"
|
||||
toctree_path = ROOT / "docs" / "source" / "_toctree.yml"
|
||||
|
||||
assert readme_path.is_symlink()
|
||||
assert readme_path.resolve() == policy_readme_path.resolve()
|
||||
assert wan_readme_path.exists()
|
||||
wan_readme = wan_readme_path.read_text()
|
||||
assert "Wan-Video/Wan2.2" in wan_readme
|
||||
assert "42bf4cfaa384bc21833865abc2f9e6c0e67233dc" in wan_readme
|
||||
assert policy_readme_path.exists()
|
||||
assert guide_path.exists()
|
||||
assert "local: fastwam" in toctree_path.read_text()
|
||||
|
||||
|
||||
def test_wan_backbone_code_is_isolated_from_lerobot_adapter():
|
||||
wan_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan"
|
||||
|
||||
assert (wan_dir / "modules" / "attention.py").exists()
|
||||
assert (wan_dir / "modules" / "model.py").exists()
|
||||
assert (wan_dir / "modules" / "t5.py").exists()
|
||||
assert (wan_dir / "modules" / "tokenizers.py").exists()
|
||||
assert (wan_dir / "modules" / "vae2_1.py").exists()
|
||||
assert (wan_dir / "modules" / "vae2_2.py").exists()
|
||||
assert (wan_dir / "utils" / "fm_solvers.py").exists()
|
||||
assert (wan_dir / "utils" / "fm_solvers_unipc.py").exists()
|
||||
|
||||
assert (wan_dir.parent / "wan_video_dit.py").exists()
|
||||
assert (wan_dir.parent / "wan_adapters.py").exists()
|
||||
assert (wan_dir.parent / "wan_components.py").exists()
|
||||
assert not (wan_dir / "wan_video_dit.py").exists()
|
||||
assert not (wan_dir / "wan_adapters.py").exists()
|
||||
assert not (wan_dir / "wan_components.py").exists()
|
||||
|
||||
|
||||
def test_fastwam_text_encoder_uses_upstream_wan_modules_directly():
|
||||
fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
|
||||
modular_source = (fastwam_dir / "modular_fastwam.py").read_text()
|
||||
components_source = (fastwam_dir / "wan_components.py").read_text()
|
||||
|
||||
assert not (fastwam_dir / "wan_video_text_encoder.py").exists()
|
||||
assert "from .wan.modules.t5 import umt5_xxl" in components_source
|
||||
assert "from .wan.modules.tokenizers import HuggingfaceTokenizer" in components_source
|
||||
assert "WAN_T5_ENCODER_KWARGS" not in components_source
|
||||
assert "wan_video_text_encoder" not in modular_source
|
||||
|
||||
|
||||
def test_fastwam_vae_reuses_upstream_wan_modules():
|
||||
fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
|
||||
vae_source = (fastwam_dir / "wan_adapters.py").read_text()
|
||||
|
||||
assert not (fastwam_dir / "wan_video_vae.py").exists()
|
||||
assert "from .wan.modules.vae2_2 import Wan2_2_VAE" in vae_source
|
||||
assert "mean = [" not in vae_source
|
||||
assert "std = [" not in vae_source
|
||||
assert "class Encoder3d_38" not in vae_source
|
||||
assert "class Decoder3d_38" not in vae_source
|
||||
assert "class VideoVAE38_" not in vae_source
|
||||
|
||||
|
||||
def test_fastwam_component_loading_uses_fixed_wan_checkpoint_layout():
|
||||
modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
|
||||
modeling_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modeling_fastwam.py").read_text()
|
||||
components_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_components.py").read_text()
|
||||
|
||||
assert "class ModelConfig" not in modular_source
|
||||
assert "def load_state_dict" not in modular_source
|
||||
assert "WAN22_MODEL_REGISTRY" not in modular_source
|
||||
assert "class ModelConfig" not in components_source
|
||||
assert "class WanComponentSource" not in components_source
|
||||
assert "def load_state_dict" not in components_source
|
||||
assert "WAN22_MODEL_REGISTRY" not in components_source
|
||||
assert "hash_model_file" not in components_source
|
||||
assert "_resolve_component_sources" not in components_source
|
||||
assert "origin_file_pattern" not in components_source
|
||||
assert "inspect.signature" not in components_source
|
||||
assert "class FastWAMWanComponentPaths" not in modeling_source
|
||||
assert "def _first_existing" not in modeling_source
|
||||
assert "def _missing_wan_component_names" not in modeling_source
|
||||
assert "WAN_T5_CHECKPOINT" in components_source
|
||||
assert "WAN_VAE_CHECKPOINT" in components_source
|
||||
assert "WAN_DIT_PATTERN" in components_source
|
||||
|
||||
|
||||
def test_fastwam_dit_reuses_upstream_wan_primitives():
|
||||
dit_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_video_dit.py").read_text()
|
||||
|
||||
assert "from .wan.modules.model import" in dit_source
|
||||
assert "WanModel" in dit_source
|
||||
for duplicated_symbol in [
|
||||
"def flash_attention(",
|
||||
"def sinusoidal_embedding_1d(",
|
||||
"def rope_apply(",
|
||||
"def unpatchify(",
|
||||
"def _dense_video_freqs(",
|
||||
"class RMSNorm(",
|
||||
"class SelfAttention(",
|
||||
"class CrossAttention(",
|
||||
"class Head(",
|
||||
]:
|
||||
assert duplicated_symbol not in dit_source
|
||||
|
||||
|
||||
def test_fastwam_inference_schedule_reuses_upstream_wan_sigmas():
|
||||
modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
|
||||
|
||||
assert "def _get_wan_sampling_sigmas" in modular_source
|
||||
assert "from .wan.utils.fm_solvers import get_sampling_sigmas" in modular_source
|
||||
assert "_get_wan_sampling_sigmas(num_inference_steps, shift)" in modular_source
|
||||
|
||||
|
||||
def test_policy_config_rejects_missing_required_image_and_action_features():
|
||||
with pytest.raises(ValueError, match="image feature"):
|
||||
FastWAMConfig(
|
||||
input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="action"):
|
||||
FastWAMConfig(
|
||||
output_features={"not_action": PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
|
||||
)
|
||||
|
||||
|
||||
def test_policy_init_calls_validate_features_even_for_prebuilt_configs(monkeypatch):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
calls = []
|
||||
|
||||
def record_validate_features():
|
||||
calls.append("called")
|
||||
|
||||
monkeypatch.setattr(cfg, "validate_features", record_validate_features)
|
||||
monkeypatch.setattr(
|
||||
FastWAMPolicy,
|
||||
"_build_core_model",
|
||||
lambda self, config: nn.Linear(1, 1),
|
||||
)
|
||||
FastWAMPolicy(cfg)
|
||||
|
||||
assert calls == ["called"]
|
||||
|
||||
|
||||
def test_required_policy_entrypoints_exist_with_discoverable_names():
|
||||
assert FastWAMPolicy.config_class is FastWAMConfig
|
||||
assert FastWAMPolicy.name == "fastwam"
|
||||
assert callable(FastWAMPolicy.reset)
|
||||
assert callable(FastWAMPolicy.get_optim_params)
|
||||
assert callable(FastWAMPolicy.predict_action_chunk)
|
||||
assert callable(FastWAMPolicy.select_action)
|
||||
assert callable(FastWAMPolicy.forward)
|
||||
assert callable(make_fastwam_pre_post_processors)
|
||||
assert make_fastwam_pre_post_processors.__name__ == "make_fastwam_pre_post_processors"
|
||||
|
||||
|
||||
def test_policy_constructor_and_forward_match_byo_template_contract():
|
||||
init_signature = inspect.signature(FastWAMPolicy.__init__)
|
||||
|
||||
assert "dataset_stats" in init_signature.parameters
|
||||
assert "core_model" not in init_signature.parameters
|
||||
assert typing.get_type_hints(FastWAMPolicy.forward)["return"] == dict[str, torch.Tensor]
|
||||
|
||||
|
||||
def test_saved_config_round_trips_policy_features(tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448))
|
||||
cfg.save_pretrained(tmp_path)
|
||||
|
||||
loaded = FastWAMConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.type == "fastwam"
|
||||
assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL
|
||||
assert loaded.action_feature.shape == (7,)
|
||||
assert loaded.robot_state_feature.shape == (8,)
|
||||
|
||||
|
||||
def test_config_from_pretrained_ignores_unknown_fields(tmp_path):
|
||||
cfg = FastWAMConfig()
|
||||
cfg.save_pretrained(tmp_path)
|
||||
config_path = tmp_path / "config.json"
|
||||
payload = config_path.read_text()
|
||||
payload = payload.replace(
|
||||
'"torch_dtype": "bfloat16"',
|
||||
'"torch_dtype": "bfloat16",\n "unknown_fastwam_field": true',
|
||||
)
|
||||
config_path.write_text(payload)
|
||||
|
||||
loaded = FastWAMConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.type == "fastwam"
|
||||
assert not hasattr(loaded, "unknown_fastwam_field")
|
||||
|
||||
|
||||
def test_config_from_pretrained_does_not_use_non_wan22_tokenizer_repo_id(tmp_path):
|
||||
cfg = FastWAMConfig()
|
||||
cfg.save_pretrained(tmp_path)
|
||||
config_path = tmp_path / "config.json"
|
||||
payload = config_path.read_text()
|
||||
payload = payload.replace(
|
||||
'"tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B"',
|
||||
'"tokenizer_model_id": "somebody/old-tokenizer"',
|
||||
)
|
||||
config_path.write_text(payload)
|
||||
|
||||
loaded = FastWAMConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.tokenizer_model_id == "Wan-AI/Wan2.2-TI2V-5B"
|
||||
@@ -1,89 +0,0 @@
|
||||
#!/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.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
|
||||
|
||||
|
||||
def test_fastwam_is_registered_in_policy_factory():
|
||||
from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
|
||||
from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
|
||||
|
||||
cfg = make_policy_config("fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
|
||||
assert isinstance(cfg, FastWAMConfig)
|
||||
assert cfg.type == "fastwam"
|
||||
assert get_policy_class("fastwam") is FastWAMPolicy
|
||||
|
||||
|
||||
def test_fastwam_pre_post_processors_are_available():
|
||||
cfg = make_policy_config("fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
|
||||
preprocessor, postprocessor = make_pre_post_processors(cfg)
|
||||
|
||||
assert preprocessor.name == "policy_preprocessor"
|
||||
assert postprocessor.name == "policy_postprocessor"
|
||||
|
||||
|
||||
def test_fastwam_postprocessor_only_adds_action_inversion_when_configured():
|
||||
from lerobot.policies.fastwam.processor_fastwam import (
|
||||
FastWAMActionInversionProcessorStep,
|
||||
FastWAMActionToggleProcessorStep,
|
||||
)
|
||||
|
||||
default_cfg = make_policy_config(
|
||||
"fastwam", action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2
|
||||
)
|
||||
_, default_postprocessor = make_pre_post_processors(default_cfg)
|
||||
|
||||
assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in default_postprocessor.steps)
|
||||
assert not any(
|
||||
isinstance(step, FastWAMActionInversionProcessorStep) for step in default_postprocessor.steps
|
||||
)
|
||||
|
||||
inverted_cfg = make_policy_config(
|
||||
"fastwam",
|
||||
action_dim=3,
|
||||
proprio_dim=2,
|
||||
action_horizon=4,
|
||||
n_action_steps=2,
|
||||
toggle_action_dimensions=[],
|
||||
invert_dimensions=[-1],
|
||||
)
|
||||
_, inverted_postprocessor = make_pre_post_processors(inverted_cfg)
|
||||
|
||||
assert any(isinstance(step, FastWAMActionInversionProcessorStep) for step in inverted_postprocessor.steps)
|
||||
|
||||
|
||||
def test_fastwam_action_inversion_processor_flips_configured_dimensions():
|
||||
from lerobot.policies.fastwam.processor_fastwam import FastWAMActionInversionProcessorStep
|
||||
|
||||
processor = FastWAMActionInversionProcessorStep(invert_dimensions=[0, -1])
|
||||
action = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
|
||||
|
||||
processed = processor.action(action)
|
||||
|
||||
assert torch.equal(processed, torch.tensor([[-1.0, 2.0, -3.0], [-4.0, 5.0, -6.0]]))
|
||||
assert torch.equal(action, torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]))
|
||||
|
||||
|
||||
def test_fastwam_rejects_non_wan22_hub_model_ids():
|
||||
from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
|
||||
|
||||
with pytest.raises(ValueError, match="model_id"):
|
||||
FastWAMConfig(model_id="somebody/other-model")
|
||||
@@ -14,21 +14,26 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_model
|
||||
from torch import nn
|
||||
|
||||
from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
|
||||
from lerobot.policies import FastWAMConfig, get_policy_class, make_policy_config, make_pre_post_processors
|
||||
from lerobot.policies.fastwam import modeling_fastwam
|
||||
from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
|
||||
from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
|
||||
from lerobot.policies.fastwam.modular_fastwam import ActionDiT, MoT
|
||||
from lerobot.policies.fastwam.wan_video_dit import (
|
||||
FastWAMAttentionBlock,
|
||||
WanVideoDiT,
|
||||
fastwam_masked_attention,
|
||||
precompute_freqs_cis,
|
||||
from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy, resolve_wan_component_paths
|
||||
from lerobot.policies.fastwam.processor_fastwam import FastWAMActionToggleProcessorStep
|
||||
from lerobot.policies.fastwam.wan_components import (
|
||||
WAN_DIT_PATTERN,
|
||||
WAN_T5_CHECKPOINT,
|
||||
WAN_T5_TOKENIZER,
|
||||
WAN_VAE_CHECKPOINT,
|
||||
resolve_wan_checkpoint_paths,
|
||||
)
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
|
||||
class FakeFastWAMCore(nn.Module):
|
||||
@@ -39,293 +44,117 @@ class FakeFastWAMCore(nn.Module):
|
||||
def training_loss(self, sample):
|
||||
assert sample["video"].ndim == 5
|
||||
assert sample["context"].ndim == 3
|
||||
return sample["action"].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0}
|
||||
return sample[ACTION].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0}
|
||||
|
||||
def infer_action(self, **kwargs):
|
||||
horizon = kwargs["action_horizon"]
|
||||
return {"action": torch.ones(horizon, 3)}
|
||||
return {"action": torch.ones(1, kwargs["action_horizon"], 3)}
|
||||
|
||||
|
||||
def _patch_core_builder(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
FastWAMPolicy,
|
||||
"_build_core_model",
|
||||
lambda self, config: FakeFastWAMCore(),
|
||||
)
|
||||
|
||||
|
||||
def test_action_attention_block_supports_mot_attention_dim_larger_than_hidden_dim():
|
||||
block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32)
|
||||
x = torch.zeros(1, 2, 16)
|
||||
context = torch.zeros(1, 3, 16)
|
||||
t_mod = torch.zeros(1, 6, 16)
|
||||
freqs = precompute_freqs_cis(8, end=2).view(2, 1, -1)
|
||||
|
||||
output = block(x, context, t_mod, freqs)
|
||||
|
||||
assert output.shape == x.shape
|
||||
assert block.self_attn.q.out_features == 32
|
||||
assert block.self_attn.o.out_features == 16
|
||||
|
||||
|
||||
def test_fastwam_masked_attention_accepts_rope_float32_qk_with_bfloat16_values():
|
||||
q = torch.zeros(1, 2, 32, dtype=torch.float32)
|
||||
k = torch.zeros(1, 2, 32, dtype=torch.float32)
|
||||
v = torch.zeros(1, 2, 32, dtype=torch.bfloat16)
|
||||
|
||||
out = fastwam_masked_attention(q=q, k=k, v=v, num_heads=4)
|
||||
|
||||
assert out.dtype == torch.float32
|
||||
assert out.shape == v.shape
|
||||
|
||||
|
||||
def test_fastwam_masked_attention_runs_fp32_when_cache_promotes_keys():
|
||||
q = torch.zeros(1, 2, 32, dtype=torch.bfloat16)
|
||||
k = torch.zeros(1, 4, 32, dtype=torch.float32)
|
||||
v = torch.zeros(1, 4, 32, dtype=torch.bfloat16)
|
||||
mask = torch.ones(2, 4, dtype=torch.bool)
|
||||
|
||||
out = fastwam_masked_attention(q=q, k=k, v=v, num_heads=4, ctx_mask=mask)
|
||||
|
||||
assert out.dtype == torch.float32
|
||||
assert out.shape == q.shape
|
||||
|
||||
|
||||
def test_attention_post_projection_casts_fp32_attention_to_block_dtype():
|
||||
block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
|
||||
dtype=torch.bfloat16
|
||||
)
|
||||
residual = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
|
||||
mixed_attn = torch.zeros(1, 2, 32, dtype=torch.float32)
|
||||
gate_msa = torch.ones(1, 16, dtype=torch.bfloat16)
|
||||
shift_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
scale_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
gate_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
|
||||
out = MoT._apply_expert_post_block(
|
||||
block=block,
|
||||
residual_x=residual,
|
||||
mixed_attn_out=mixed_attn,
|
||||
gate_msa=gate_msa,
|
||||
shift_mlp=shift_mlp,
|
||||
scale_mlp=scale_mlp,
|
||||
gate_mlp=gate_mlp,
|
||||
context_payload=None,
|
||||
)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
assert out.shape == residual.shape
|
||||
|
||||
|
||||
def test_attention_cross_projection_casts_fp32_attention_to_block_dtype():
|
||||
block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
|
||||
dtype=torch.bfloat16
|
||||
)
|
||||
x = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
|
||||
context = torch.zeros(1, 3, 16, dtype=torch.bfloat16)
|
||||
|
||||
out = block.apply_cross_attention(x, context)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
assert out.shape == x.shape
|
||||
|
||||
|
||||
def test_attention_norm3_handles_bfloat16_affine_weights():
|
||||
block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
|
||||
dtype=torch.bfloat16
|
||||
)
|
||||
x = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
|
||||
|
||||
out = block.apply_norm3(x)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
assert out.shape == x.shape
|
||||
|
||||
|
||||
def test_attention_post_block_handles_bfloat16_cross_attention_norm():
|
||||
block = FastWAMAttentionBlock(hidden_dim=16, attn_head_dim=8, num_heads=4, ffn_dim=32).to(
|
||||
dtype=torch.bfloat16
|
||||
)
|
||||
residual = torch.zeros(1, 2, 16, dtype=torch.bfloat16)
|
||||
mixed_attn = torch.zeros(1, 2, 32, dtype=torch.float32)
|
||||
gate_msa = torch.ones(1, 16, dtype=torch.bfloat16)
|
||||
shift_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
scale_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
gate_mlp = torch.zeros(1, 16, dtype=torch.bfloat16)
|
||||
context_payload = {"context": torch.zeros(1, 3, 16, dtype=torch.bfloat16), "mask": None}
|
||||
|
||||
out = MoT._apply_expert_post_block(
|
||||
block=block,
|
||||
residual_x=residual,
|
||||
mixed_attn_out=mixed_attn,
|
||||
gate_msa=gate_msa,
|
||||
shift_mlp=shift_mlp,
|
||||
scale_mlp=scale_mlp,
|
||||
gate_mlp=gate_mlp,
|
||||
context_payload=context_payload,
|
||||
)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
assert out.shape == residual.shape
|
||||
|
||||
|
||||
def test_video_dit_pre_dit_casts_double_latents_to_model_dtype():
|
||||
model = WanVideoDiT(
|
||||
hidden_dim=4,
|
||||
in_dim=48,
|
||||
ffn_dim=8,
|
||||
out_dim=48,
|
||||
text_dim=6,
|
||||
freq_dim=4,
|
||||
eps=1e-6,
|
||||
patch_size=(1, 2, 2),
|
||||
num_heads=1,
|
||||
attn_head_dim=4,
|
||||
num_layers=0,
|
||||
seperated_timestep=True,
|
||||
fuse_vae_embedding_in_latents=True,
|
||||
video_attention_mask_mode="first_frame_causal",
|
||||
).to(dtype=torch.bfloat16)
|
||||
|
||||
state = model.pre_dit(
|
||||
x=torch.zeros(1, 48, 1, 2, 2, dtype=torch.float64),
|
||||
timestep=torch.zeros(1, dtype=torch.float64),
|
||||
context=torch.zeros(1, 2, 6, dtype=torch.float64),
|
||||
fuse_vae_embedding_in_latents=True,
|
||||
)
|
||||
|
||||
assert state["tokens"].dtype == torch.bfloat16
|
||||
assert state["context"].dtype == torch.bfloat16
|
||||
assert state["t_mod"].dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_action_dit_pre_dit_casts_double_inputs_to_model_dtype():
|
||||
model = ActionDiT(
|
||||
hidden_dim=16,
|
||||
def test_fastwam_is_registered_and_publicly_exported():
|
||||
cfg = make_policy_config(
|
||||
"fastwam",
|
||||
action_dim=3,
|
||||
ffn_dim=32,
|
||||
text_dim=6,
|
||||
freq_dim=4,
|
||||
eps=1e-6,
|
||||
num_heads=4,
|
||||
attn_head_dim=8,
|
||||
num_layers=0,
|
||||
).to(dtype=torch.bfloat16)
|
||||
|
||||
state = model.pre_dit(
|
||||
action_tokens=torch.zeros(1, 2, 3, dtype=torch.float64),
|
||||
timestep=torch.zeros(1, dtype=torch.float64),
|
||||
context=torch.zeros(1, 2, 6, dtype=torch.float64),
|
||||
proprio_dim=2,
|
||||
action_horizon=4,
|
||||
n_action_steps=2,
|
||||
base_model_id=None,
|
||||
)
|
||||
|
||||
assert state["tokens"].dtype == torch.bfloat16
|
||||
assert state["context"].dtype == torch.bfloat16
|
||||
assert state["t_mod"].dtype == torch.bfloat16
|
||||
assert isinstance(cfg, FastWAMConfig)
|
||||
assert cfg.type == "fastwam"
|
||||
assert get_policy_class("fastwam") is FastWAMPolicy
|
||||
|
||||
|
||||
def test_forward_adapts_lerobot_batch_to_fastwam_sample(monkeypatch):
|
||||
_patch_core_builder(monkeypatch)
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
batch = {
|
||||
"observation.images.image": torch.zeros(1, 3, 16, 16),
|
||||
"observation.state": torch.zeros(1, 2),
|
||||
"action": torch.zeros(1, 4, 3),
|
||||
"context": torch.zeros(1, 5, 4096),
|
||||
"context_mask": torch.ones(1, 5, dtype=torch.bool),
|
||||
}
|
||||
def test_config_validates_features_model_ids_and_saved_auto_route(tmp_path):
|
||||
cfg = FastWAMConfig()
|
||||
cfg.save_pretrained(tmp_path)
|
||||
saved = json.loads((tmp_path / "config.json").read_text())
|
||||
|
||||
output = policy.forward(batch)
|
||||
|
||||
assert set(output) == {"loss", "loss_action"}
|
||||
assert output["loss"].item() == 1.0
|
||||
assert output["loss_action"].item() == 1.0
|
||||
assert saved["pretrained_path"] is None
|
||||
assert cfg.image_features["observation.images.image"].type == FeatureType.VISUAL
|
||||
assert cfg.action_feature.shape == (7,)
|
||||
assert cfg.robot_state_feature.shape == (8,)
|
||||
with pytest.raises(ValueError, match="image feature"):
|
||||
FastWAMConfig(input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))})
|
||||
with pytest.raises(ValueError, match="tokenizer_model_id"):
|
||||
FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer")
|
||||
|
||||
|
||||
def test_get_optim_params_returns_lerobot_optimizer_dict(monkeypatch):
|
||||
_patch_core_builder(monkeypatch)
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
|
||||
optim_params = policy.get_optim_params()
|
||||
|
||||
assert isinstance(optim_params, dict)
|
||||
assert set(optim_params) == {"params"}
|
||||
assert list(optim_params["params"])
|
||||
|
||||
|
||||
def test_select_action_uses_action_queue(monkeypatch):
|
||||
_patch_core_builder(monkeypatch)
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
batch = {
|
||||
"input_image": torch.zeros(1, 3, 16, 16),
|
||||
"observation.state": torch.zeros(1, 2),
|
||||
"context": torch.zeros(1, 5, 4096),
|
||||
"context_mask": torch.ones(1, 5, dtype=torch.bool),
|
||||
}
|
||||
|
||||
first = policy.select_action(batch)
|
||||
second = policy.select_action(batch)
|
||||
|
||||
assert first.shape == (1, 3)
|
||||
assert second.shape == (1, 3)
|
||||
|
||||
|
||||
def test_predict_action_prepares_lerobot_libero_observation(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
class CapturingCore(FakeFastWAMCore):
|
||||
def infer_action(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return {"action": torch.ones(1, 4, 3)}
|
||||
|
||||
monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore())
|
||||
def test_preprocessor_normalizes_images_and_postprocessor_toggles_actions(tmp_path):
|
||||
cfg = FastWAMConfig(
|
||||
action_dim=3,
|
||||
proprio_dim=2,
|
||||
action_horizon=4,
|
||||
n_action_steps=2,
|
||||
image_size=(16, 32),
|
||||
image_size=(2, 2),
|
||||
device="cpu",
|
||||
toggle_action_dimensions=[-1],
|
||||
input_features={
|
||||
"observation.images.image": {"type": "VISUAL", "shape": (3, 16, 32)},
|
||||
"observation.state": {"type": "STATE", "shape": (2,)},
|
||||
"observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 2, 2)),
|
||||
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
|
||||
},
|
||||
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
|
||||
base_model_id=None,
|
||||
)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
batch = {
|
||||
"observation.images.image": torch.ones(1, 3, 20, 20),
|
||||
"observation.images.image2": torch.zeros(1, 3, 20, 20),
|
||||
"observation.state": torch.zeros(1, 2),
|
||||
"task": ["pick up the bowl"],
|
||||
dataset_stats = {
|
||||
"observation.images.image": {
|
||||
"mean": torch.full((3, 1, 1), 0.2),
|
||||
"std": torch.full((3, 1, 1), 0.1),
|
||||
},
|
||||
OBS_STATE: {
|
||||
"mean": torch.tensor([1.0, 3.0]),
|
||||
"std": torch.tensor([2.0, 4.0]),
|
||||
},
|
||||
ACTION: {
|
||||
"mean": torch.zeros(3),
|
||||
"std": torch.ones(3),
|
||||
},
|
||||
}
|
||||
|
||||
action = policy.predict_action_chunk(batch)
|
||||
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats)
|
||||
processed = preprocessor(
|
||||
{
|
||||
"observation.images.image": torch.tensor(
|
||||
[
|
||||
[[0.0, 0.5], [1.0, 0.5]],
|
||||
[[0.0, 0.5], [1.0, 0.5]],
|
||||
[[0.0, 0.5], [1.0, 0.5]],
|
||||
]
|
||||
),
|
||||
OBS_STATE: torch.tensor([3.0, 7.0]),
|
||||
}
|
||||
)
|
||||
preprocessor.save_pretrained(tmp_path, config_filename="policy_preprocessor.json")
|
||||
postprocessor.save_pretrained(tmp_path, config_filename="policy_postprocessor.json")
|
||||
_, loaded_postprocessor = make_pre_post_processors(cfg, pretrained_path=str(tmp_path))
|
||||
|
||||
assert action.shape == (1, 4, 3)
|
||||
assert captured["prompt"] == [cfg.prompt_template.format(task="pick up the bowl")]
|
||||
assert tuple(captured["input_image"].shape) == (1, 3, 16, 32)
|
||||
assert captured["input_image"].amin().item() == -1.0
|
||||
assert captured["input_image"].amax().item() == 1.0
|
||||
assert "num_video_frames" not in captured
|
||||
expected_image = torch.tensor(
|
||||
[[[[-1.0, 0.0], [1.0, 0.0]], [[-1.0, 0.0], [1.0, 0.0]], [[-1.0, 0.0], [1.0, 0.0]]]]
|
||||
)
|
||||
assert preprocessor.name == "policy_preprocessor"
|
||||
assert postprocessor.name == "policy_postprocessor"
|
||||
assert torch.allclose(processed["observation.images.image"], expected_image)
|
||||
assert torch.allclose(processed[OBS_STATE], torch.tensor([[1.0, 1.0]]))
|
||||
assert torch.equal(dataset_stats["observation.images.image"]["mean"], torch.full((3, 1, 1), 0.2))
|
||||
assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in loaded_postprocessor.steps)
|
||||
assert torch.equal(
|
||||
loaded_postprocessor(torch.tensor([[0.25, 0.5, 1.0]])), torch.tensor([[0.25, 0.5, -1.0]])
|
||||
)
|
||||
|
||||
|
||||
def test_predict_action_splits_parallel_eval_batch_into_single_infer_calls(monkeypatch):
|
||||
def test_policy_forward_and_predict_action_adapt_lerobot_batches(monkeypatch):
|
||||
captured = []
|
||||
|
||||
class CapturingCore(FakeFastWAMCore):
|
||||
def infer_action(self, **kwargs):
|
||||
captured.append(
|
||||
{
|
||||
"input_image_shape": tuple(kwargs["input_image"].shape),
|
||||
"input_image_sum": float(kwargs["input_image"].sum()),
|
||||
"image_shape": tuple(kwargs["input_image"].shape),
|
||||
"proprio_shape": tuple(kwargs["proprio"].shape),
|
||||
"proprio_sum": float(kwargs["proprio"].sum()),
|
||||
"prompt": kwargs["prompt"],
|
||||
}
|
||||
)
|
||||
action = torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))
|
||||
return {"action": action}
|
||||
return {"action": torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))}
|
||||
|
||||
monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore())
|
||||
cfg = FastWAMConfig(
|
||||
@@ -335,42 +164,54 @@ def test_predict_action_splits_parallel_eval_batch_into_single_infer_calls(monke
|
||||
n_action_steps=2,
|
||||
image_size=(16, 16),
|
||||
input_features={
|
||||
"observation.images.image": {"type": "VISUAL", "shape": (3, 16, 16)},
|
||||
"observation.state": {"type": "STATE", "shape": (2,)},
|
||||
"observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
|
||||
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
|
||||
},
|
||||
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
|
||||
base_model_id=None,
|
||||
)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
batch = {
|
||||
"observation.images.image": torch.stack(
|
||||
[
|
||||
torch.zeros(3, 16, 16),
|
||||
torch.ones(3, 16, 16),
|
||||
torch.full((3, 16, 16), 2.0),
|
||||
]
|
||||
),
|
||||
"observation.state": torch.tensor([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]),
|
||||
"task": ["task 0", "task 1", "task 2"],
|
||||
}
|
||||
with pytest.warns(RuntimeWarning, match="does not load pretrained FastWAM weights"):
|
||||
policy = FastWAMPolicy(cfg)
|
||||
|
||||
action = policy.predict_action_chunk(batch)
|
||||
output = policy.forward(
|
||||
{
|
||||
"observation.images.image": torch.zeros(1, 3, 16, 16),
|
||||
OBS_STATE: torch.zeros(1, 2),
|
||||
ACTION: torch.zeros(1, 4, 3),
|
||||
"context": torch.zeros(1, 5, 4096),
|
||||
"context_mask": torch.ones(1, 5, dtype=torch.bool),
|
||||
}
|
||||
)
|
||||
action = policy.predict_action_chunk(
|
||||
{
|
||||
"observation.images.image": torch.stack(
|
||||
[
|
||||
torch.zeros(3, 16, 16),
|
||||
torch.ones(3, 16, 16),
|
||||
]
|
||||
),
|
||||
OBS_STATE: torch.tensor([[0.0, 1.0], [2.0, 3.0]]),
|
||||
"task": ["task 0", "task 1"],
|
||||
}
|
||||
)
|
||||
|
||||
assert action.shape == (3, 4, 3)
|
||||
assert action[:, 0, 0].tolist() == [1.0, 2.0, 3.0]
|
||||
assert len(captured) == 3
|
||||
assert [item["input_image_shape"] for item in captured] == [(1, 3, 16, 16)] * 3
|
||||
assert [item["proprio_shape"] for item in captured] == [(1, 2)] * 3
|
||||
assert output["loss"].item() == 1.0
|
||||
assert output["loss_action"].item() == 1.0
|
||||
assert action.shape == (2, 4, 3)
|
||||
assert action[:, 0, 0].tolist() == [1.0, 2.0]
|
||||
assert [item["image_shape"] for item in captured] == [(1, 3, 16, 16), (1, 3, 16, 16)]
|
||||
assert [item["proprio_shape"] for item in captured] == [(1, 2), (1, 2)]
|
||||
assert [item["prompt"] for item in captured] == [
|
||||
cfg.prompt_template.format(task="task 0"),
|
||||
cfg.prompt_template.format(task="task 1"),
|
||||
cfg.prompt_template.format(task="task 2"),
|
||||
]
|
||||
|
||||
|
||||
def test_from_pretrained_does_not_initialize_wan_backbone(monkeypatch, tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
def test_from_pretrained_loads_weights_without_initializing_wan_backbone(monkeypatch, tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2, base_model_id=None)
|
||||
cfg.save_pretrained(tmp_path)
|
||||
_patch_core_builder(monkeypatch)
|
||||
reference_policy = FastWAMPolicy(cfg)
|
||||
monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: FakeFastWAMCore())
|
||||
reference_policy = FastWAMPolicy(cfg, _suppress_base_init_warning=True)
|
||||
save_model(reference_policy, str(tmp_path / "model.safetensors"))
|
||||
|
||||
def fail_if_wan_pretrained_is_loaded(*args, **kwargs):
|
||||
@@ -399,52 +240,13 @@ def test_from_pretrained_does_not_initialize_wan_backbone(monkeypatch, tmp_path)
|
||||
assert loaded_components_from == [tmp_path]
|
||||
|
||||
|
||||
def test_from_pretrained_resolves_hub_repo_to_snapshot_before_loading_sidecars(monkeypatch, tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
cfg.save_pretrained(tmp_path)
|
||||
snapshot_calls = []
|
||||
|
||||
def fake_snapshot_download(**kwargs):
|
||||
snapshot_calls.append(kwargs)
|
||||
return str(tmp_path)
|
||||
|
||||
@classmethod
|
||||
def fake_base_from_pretrained(cls, pretrained_name_or_path, *, config=None, **kwargs):
|
||||
assert pretrained_name_or_path == tmp_path
|
||||
assert kwargs.pop("_skip_wan_init") is True
|
||||
assert kwargs["strict"] is False
|
||||
return cls(config, _skip_wan_init=True)
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(PreTrainedPolicy, "from_pretrained", fake_base_from_pretrained)
|
||||
monkeypatch.setattr(
|
||||
modeling_fastwam,
|
||||
"_build_core_model_from_architecture",
|
||||
lambda config: FakeFastWAMCore(),
|
||||
raising=False,
|
||||
)
|
||||
loaded_components_from = []
|
||||
monkeypatch.setattr(
|
||||
FastWAMPolicy,
|
||||
"load_wan_components_from_pretrained",
|
||||
lambda self, path: loaded_components_from.append(path),
|
||||
)
|
||||
|
||||
FastWAMPolicy.from_pretrained("org/fastwam", strict=False, local_files_only=True, revision="main")
|
||||
|
||||
assert snapshot_calls[0]["repo_id"] == "org/fastwam"
|
||||
assert snapshot_calls[0]["local_files_only"] is True
|
||||
assert snapshot_calls[0]["revision"] == "main"
|
||||
assert loaded_components_from == [tmp_path]
|
||||
|
||||
|
||||
def test_save_pretrained_copies_wan_components(monkeypatch, tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
|
||||
def test_save_pretrained_copies_required_wan_sidecars(monkeypatch, tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2, base_model_id=None)
|
||||
source = tmp_path / "source"
|
||||
tokenizer = source / "google" / "umt5-xxl"
|
||||
tokenizer = source / WAN_T5_TOKENIZER
|
||||
tokenizer.mkdir(parents=True)
|
||||
vae = source / "Wan2.2_VAE.pth"
|
||||
text_encoder = source / "models_t5_umt5-xxl-enc-bf16.pth"
|
||||
vae = source / WAN_VAE_CHECKPOINT
|
||||
text_encoder = source / WAN_T5_CHECKPOINT
|
||||
tokenizer_file = tokenizer / "tokenizer.json"
|
||||
vae.write_bytes(b"vae")
|
||||
text_encoder.write_bytes(b"text")
|
||||
@@ -456,12 +258,47 @@ def test_save_pretrained_copies_wan_components(monkeypatch, tmp_path):
|
||||
"tokenizer": str(tokenizer),
|
||||
}
|
||||
monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: core)
|
||||
policy = FastWAMPolicy(cfg)
|
||||
policy = FastWAMPolicy(cfg, _suppress_base_init_warning=True)
|
||||
|
||||
save_dir = tmp_path / "saved"
|
||||
policy.save_pretrained(save_dir)
|
||||
|
||||
assert (save_dir / "model.safetensors").is_file()
|
||||
assert (save_dir / "Wan2.2_VAE.pth").read_bytes() == b"vae"
|
||||
assert (save_dir / "models_t5_umt5-xxl-enc-bf16.pth").read_bytes() == b"text"
|
||||
assert (save_dir / "google" / "umt5-xxl" / "tokenizer.json").read_text() == "{}"
|
||||
assert (save_dir / WAN_VAE_CHECKPOINT).read_bytes() == b"vae"
|
||||
assert (save_dir / WAN_T5_CHECKPOINT).read_bytes() == b"text"
|
||||
assert (save_dir / WAN_T5_TOKENIZER / "tokenizer.json").read_text() == "{}"
|
||||
|
||||
|
||||
def test_wan_component_resolution_uses_fixed_safetensors_layout(tmp_path):
|
||||
tokenizer = tmp_path / WAN_T5_TOKENIZER
|
||||
tokenizer.mkdir(parents=True)
|
||||
(tmp_path / WAN_VAE_CHECKPOINT).touch()
|
||||
(tmp_path / WAN_T5_CHECKPOINT).touch()
|
||||
(tmp_path / "diffusion_pytorch_model-00001-of-00001.safetensors").touch()
|
||||
(tokenizer / "tokenizer.json").touch()
|
||||
|
||||
paths = resolve_wan_checkpoint_paths(tmp_path)
|
||||
sidecar_paths = resolve_wan_component_paths(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 sidecar_paths.dit == []
|
||||
assert WAN_DIT_PATTERN == "diffusion_pytorch_model*.safetensors"
|
||||
|
||||
(tmp_path / WAN_T5_CHECKPOINT).unlink()
|
||||
with pytest.raises(FileNotFoundError, match="text encoder"):
|
||||
resolve_wan_checkpoint_paths(tmp_path)
|
||||
|
||||
|
||||
def test_pretrained_config_round_trips_fastwam_features(tmp_path):
|
||||
cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448), base_model_id=None)
|
||||
cfg.save_pretrained(tmp_path)
|
||||
|
||||
loaded = PreTrainedConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.type == "fastwam"
|
||||
assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL
|
||||
assert loaded.action_feature.shape == (7,)
|
||||
assert loaded.robot_state_feature.shape == (8,)
|
||||
|
||||
@@ -1,162 +0,0 @@
|
||||
#!/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
|
||||
Reference in New Issue
Block a user