Problems when directly running inference

[Charles-ux-bit]

I use command to run inference:
python scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC (note that a space in "--" and "use_latest_checkpoints")

but the command return error:

[Errno 2] No such file or directory: C:\Users\M\Desktop\TinyWorld\results\20251117_132228_sonic\video_tokenizer\checkpoints\sonic\sonic_video_tokenizer_step_27500_2025_09_17_06_20_26.pth\model_state_dict.pt

the path "C:\Users\M\Desktop\TinyWorld\results\20251117_132228_sonic\video_tokenizer\checkpoints\sonic\sonic_video_tokenizer_step_27500_2025_09_17_06_20_26.pth" exists, but there is indeed no model_state_dict.pt file.

Is there anything wrong? Thank you.

[AlmondGod]

did you download the model from huggingface?

[erenovic]

I think the model checkpoints that are accessed by the repo do not conform the model loading convention. Changing the code into this solved the issue for me but by no means, I am not sure if it is compatible with rest of the code for me:

def load_videotokenizer_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False):
    """Instantiate VideoTokenizer from a checkpoint's saved config and load weights."""
    import torch
    from models.video_tokenizer import VideoTokenizer
    model_sd = torch.load(Path(checkpoint_path), map_location='cpu', weights_only=True)  # Changed
    state_cfg = torch.load(Path(checkpoint_path), map_location='cpu', weights_only=False)  # Changed
    cfg = state_cfg.get('config', {}) or {}
    frame_size = cfg.get('frame_size', 128)
    kwargs = {
        'frame_size': (frame_size, frame_size),
        'patch_size': cfg.get('patch_size', 8),
        'embed_dim': cfg.get('embed_dim', 128),
        'num_heads': cfg.get('num_heads', 8),
        'hidden_dim': cfg.get('hidden_dim', 256),
        'num_blocks': cfg.get('num_blocks', 4),
        'latent_dim': cfg.get('latent_dim', 6),
        'num_bins': cfg.get('num_bins', 4),
    }
    if model is None:
        model = VideoTokenizer(**kwargs)
    set_model_state_dict(
        model=model,
        model_state_dict=model_sd["model"] if "model" in model_sd else model_sd,  #Changed
        options=StateDictOptions(
            full_state_dict=True,
            broadcast_from_rank0=is_distributed,
        ),
    )
    model = model.to(device)
    return model, state_cfg

state_cfg is not used once it is loaded (see load_models in inference_utils.py) so don't worry about that part, at least for inference.

You need to repeat the same modifications for load_latent_actions_from_checkpoint and load_dynamics_from_checkpoint, too. Also, the results I got were quite bad but I doubt that this is related to model loading. I suspect that there might be a CUDA - mps mismatch as I am using CUDA.

[pedrohenriquebr]

I tried the same fixing but I got another errors lol as follow

(tinyworlds) C:\Users\HP\OneDrive\Documentos\dev\tinyworlds>uv run scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC
W0105 23:13:27.726000 5412 .venv\Lib\site-packages\torch\distributed\elastic\multiprocessing\redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.
Using video_tokenizer checkpoint: C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\results\20260105_225540_sonic\video_tokenizer\checkpoints\sonic\sonic_video_tokenizer_step_27500_2025_09_17_06_20_26.pth
Using latent_actions checkpoint: C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\results\20260105_225540_sonic\actions\checkpoints\sonic\sonic_latent_actions_step_2500_2025_09_17_06_50_59.pth
Using dynamics checkpoint: C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\results\20260105_225540_sonic\dynamics\checkpoints\sonic\sonic_dynamics_step_97500_2025_09_18_11_25_59.pth
Traceback (most recent call last):
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\scripts\run_inference.py", line 145, in
main()
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\scripts\run_inference.py", line 59, in main
video_tokenizer, latent_action_model, dynamics_model = load_models(args.video_tokenizer_path, args.latent_actions_path, args.dynamics_path, args.device, use_actions=use_latent_actions)
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\utils\inference_utils.py", line 18, in load_models
dynamics_model, _dyn_ckpt = load_dynamics_from_checkpoint(dynamics_path, device)
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds\utils\utils.py", line 304, in load_dynamics_from_checkpoint
set_model_state_dict(
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds.venv\lib\site-packages\torch\distributed\checkpoint\state_dict.py", line 1271, in set_model_state_dict
return _load_model_state_dict(model, model_state_dict, info)
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds.venv\lib\site-packages\torch\utils_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds.venv\lib\site-packages\torch\distributed\checkpoint\state_dict.py", line 599, in _load_model_state_dict
_state_dict_fn(model, "load_state_dict")(
File "C:\Users\HP\OneDrive\Documentos\dev\tinyworlds.venv\lib\site-packages\torch\nn\modules\module.py", line 2629, in load_state_dict
raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for DynamicsModel:
size mismatch for transformer.blocks.0.spatial_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.0.temporal_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.0.ffn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.1.spatial_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.1.temporal_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.1.ffn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.2.spatial_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.2.temporal_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.2.ffn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.3.spatial_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.3.temporal_attn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).
size mismatch for transformer.blocks.3.ffn.norm.to_gamma_beta.1.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([256, 3]).

[GYXXXXXX]

You need to repeat the same modifications for load_latent_actions_from_checkpoint and load_dynamics_from_checkpoint, too. Also, the results I got were quite bad but I doubt that this is related to model loading. I suspect that there might be a CUDA - mps mismatch as I am using CUDA

@erenovic
Thank you for your modification. However, I encountered the same issue after the modification, and the results I got using CUDA were also terrible. Could you please tell me how you solved it in the end?

[AlmondGod]

hmm this looks like a mismatched dynamics model-tokenizer as far as I can tell

Prior error you need to change action size in config perhaps

[GYXXXXXX]

hmm this looks like a mismatched dynamics model-tokenizer as far as I can tell

Prior error you need to change action size in config perhaps

Thank you very much for your reply. The inference result was performed using the pretrained model you provided. If the model is trained from scratch, should the value of n_actions be adjusted according to different datasets? Is a lower resolution associated with a smaller n_actions value?Looking forward to your response.