hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

oops

Browse Source
This commit is contained in:
lucidrains 2025-10-21 09:50:07 -07:00
parent 4a5465eeb6
commit 283d59d75a
3 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
dreamer4/dreamer4.py
View File

@ -1251,6 +1251,7 @@ class AxialSpaceTimeTransformer(Module):
final_norm = True final_norm = True
): ):
super().__init__() super().__init__()
assert depth >= time_block_every, f'depth must be at least {time_block_every}'
# hyper connections # hyper connections

2
pyproject.toml
View File

@ -1,6 +1,6 @@
[project] [project]
name = "dreamer4" name = "dreamer4"
version = "0.0.55" version = "0.0.56"
description = "Dreamer 4" description = "Dreamer 4"
authors = [ authors = [
{ name = "Phil Wang", email = "lucidrains@gmail.com" } { name = "Phil Wang", email = "lucidrains@gmail.com" }

11
tests/test_dreamer.py
View File

@ -30,8 +30,8 @@ def test_e2e(
tokenizer = VideoTokenizer( tokenizer = VideoTokenizer(
16, 16,
encoder_depth = 1, encoder_depth = 4,
decoder_depth = 1, decoder_depth = 4,
dim_latent = 16, dim_latent = 16,
patch_size = 32, patch_size = 32,
attn_dim_head = 16, attn_dim_head = 16,
@ -197,8 +197,8 @@ def test_action_with_world_model():
512, 512,
dim_latent = 32, dim_latent = 32,
patch_size = 32, patch_size = 32,
encoder_depth = 1, encoder_depth = 4,
decoder_depth = 1, decoder_depth = 4,
attn_heads = 8, attn_heads = 8,
image_height = 256, image_height = 256,
image_width = 256, image_width = 256,
@ -212,6 +212,7 @@ def test_action_with_world_model():
num_agents = 1, num_agents = 1,
video_tokenizer = tokenizer, video_tokenizer = tokenizer,
dim_latent = 32, dim_latent = 32,
depth = 4,
num_discrete_actions = 4 num_discrete_actions = 4
) )
@ -483,7 +484,7 @@ def test_cache_generate():
max_steps = 64, max_steps = 64,
num_tasks = 4, num_tasks = 4,
num_latent_tokens = 4, num_latent_tokens = 4,
depth = 1, depth = 4,
num_spatial_tokens = 1, num_spatial_tokens = 1,
pred_orig_latent = True, pred_orig_latent = True,
num_discrete_actions = 4, num_discrete_actions = 4,