Update config.py

config.py

@@ -11,16 +11,16 @@ from diffuse.models import CondUNet2D
 
 # Model definition
 model = CondUNet2D(
-        in_channels=IN_CHANNELS,
-        ch=variant["base_channels"],
-        ch_mult=variant["channel_multipliers"],
-        num_res_blocks=variant["num_res_blocks"],
-        attention_resolutions=variant["attention_resolutions"],
-        num_heads=variant["num_heads"],
-        dropout=True,
-        dropout_rate=0.1,
-        activation=nnx.swish,
-        param_dtype=jnp.float32,
-        dtype=compute_dtype,
-        rngs=nnx.Rngs(MODEL_SEED),
-    )
+    in_channels=IN_CHANNELS,
+    ch=32,  # base_channels
+    ch_mult=(1, 2, 2),  # channel_multipliers
+    num_res_blocks=1,
+    attention_resolutions=(2,),
+    num_heads=4,
+    dropout=True,
+    dropout_rate=0.1,
+    activation=nnx.swish,
+    param_dtype=jnp.float32,
+    rngs=nnx.Rngs(0),
+)
+