---
library_name: transformers
base_model:
- google/diffusiongemma-26B-A4B-it
---
This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from [google/diffusiongemma-26B-A4B-it](https://huggingface.co/google/diffusiongemma-26B-A4B-it).
Structural comparison with `google/diffusiongemma-26B-A4B-it`:
| Feature | Source | Tiny |
|---|---:|---:|
| Input modalities | text, image, video | same |
| Vocabulary / tied embeddings | 262,144 / yes | same |
| Text layers | 30 | 12; two 5-sliding + 1-full cycles |
| Text width / local head dimension | 2,816 / 256 | 16 / 32 |
| Full-attention head dimension | 512 | 64 |
| MoE experts / active experts | 128 / 8 | 32 / 8 |
| Vision layers / width | 27 / 1,152 | 2 / 64 |
| Canvas length and self-conditioning | 256 / present | same |
| Stored precision | BF16 | BF16 |
The encoder/decoder weight sharing and the decoder self-conditioning branch are
provided by the native `DiffusionGemmaForBlockDiffusion` implementation rather
than serialized as independent copies.
| File path | Size |
|------|------|
| model.safetensors | 12.6MB |
### Example usage:
- Transformers (✅ verified)
```python
import numpy as np
import torch
from PIL import Image
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
from transformers.video_utils import VideoMetadata
model_id = "tiny-random/diffusiongemma-moe"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=dtype).to(device)
processor = AutoProcessor.from_pretrained(model_id)
image = Image.new("RGB", (64, 64), "red")
video = np.stack([np.asarray(Image.new("RGB", (64, 64), color)) for color in ("red", "green", "blue", "yellow")])
messages = [{"role": "user", "content": [
{"type": "image", "image": image},
{"type": "video", "video": video},
{"type": "text", "text": "Describe the image and video."},
]}]
inputs = processor.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt",
do_sample_frames=False,
video_metadata=[VideoMetadata(total_num_frames=4, fps=2, frames_indices=[0, 1, 2, 3])],
).to(device)
output = model.generate(**inputs, max_new_tokens=10, max_denoising_steps=2)
print(processor.decode(output.sequences[0], skip_special_tokens=True)[-500:])
```
### Codes to create this repo (✅ verified):
Click to expand
```python
import json
from pathlib import Path
from shutil import copyfile
import torch
from huggingface_hub import hf_hub_download, list_repo_files
from transformers import AutoConfig, DiffusionGemmaForBlockDiffusion, set_seed
source_model_id = "google/diffusiongemma-26B-A4B-it"
save_folder = "/tmp/tiny-random/diffusiongemma-moe"
source_revision = "f7f5b7f5fa82ffc52addd066915886d497f5517b"
set_seed(42)
folder = Path(save_folder)
folder.mkdir(parents=True, exist_ok=True)
raw = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text())
text = raw["text_config"]
text.update(
hidden_size=16, head_dim=32, num_attention_heads=1, num_key_value_heads=1,
intermediate_size=64, moe_intermediate_size=32, num_experts=32, top_k_experts=8,
num_hidden_layers=12,
layer_types=["sliding_attention"] * 5 + ["full_attention"] + ["sliding_attention"] * 5 + ["full_attention"],
per_layer_config={"5": {"head_dim": 64, "num_key_value_heads": 1}, "11": {"head_dim": 64, "num_key_value_heads": 1}},
)
raw["vision_config"].update(
hidden_size=64, intermediate_size=128, num_hidden_layers=2,
num_attention_heads=2, num_key_value_heads=2, head_dim=32,
)
(folder / "config.json").write_text(json.dumps(raw, indent=2), encoding="utf-8")
config = AutoConfig.from_pretrained(folder)
model = DiffusionGemmaForBlockDiffusion(config).to(torch.bfloat16)
with torch.no_grad():
for name, parameter in sorted(model.named_parameters()):
if not parameter.is_floating_point():
continue
if name.endswith("norm.weight"):
parameter.fill_(1)
else:
torch.nn.init.normal_(parameter, mean=0, std=0.02)
model.save_pretrained(folder, safe_serialization=True)
assets = {"chat_template.jinja", "generation_config.json", "model_index.json", "processor_config.json", "scheduler/scheduler_config.json", "tokenizer.json", "tokenizer_config.json"}
for filename in sorted(set(list_repo_files(source_model_id, revision=source_revision)) & assets):
destination = folder / filename
destination.parent.mkdir(parents=True, exist_ok=True)
copyfile(hf_hub_download(source_model_id, filename, revision=source_revision), destination)
```
Click to expand
```text
DiffusionGemmaForBlockDiffusion(
(model): DiffusionGemmaModel(
(encoder): DiffusionGemmaEncoderModel(
(language_model): DiffusionGemmaEncoderTextModel(
(embed_tokens): DiffusionGemmaTextScaledWordEmbedding(262144, 16, padding_idx=0)
(layers): ModuleList(
(0-4): 5 x DiffusionGemmaEncoderTextLayer(
(self_attn): DiffusionGemmaEncoderTextAttention(
(q_proj): Linear(in_features=16, out_features=32, bias=False)
(k_proj): Linear(in_features=16, out_features=32, bias=False)
(v_proj): Linear(in_features=16, out_features=32, bias=False)
(o_proj): Linear(in_features=32, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(5): DiffusionGemmaEncoderTextLayer(
(self_attn): DiffusionGemmaEncoderTextAttention(
(q_proj): Linear(in_features=16, out_features=64, bias=False)
(k_proj): Linear(in_features=16, out_features=64, bias=False)
(o_proj): Linear(in_features=64, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(6-10): 5 x DiffusionGemmaEncoderTextLayer(
(self_attn): DiffusionGemmaEncoderTextAttention(
(q_proj): Linear(in_features=16, out_features=32, bias=False)
(k_proj): Linear(in_features=16, out_features=32, bias=False)
(v_proj): Linear(in_features=16, out_features=32, bias=False)
(o_proj): Linear(in_features=32, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(11): DiffusionGemmaEncoderTextLayer(
(self_attn): DiffusionGemmaEncoderTextAttention(
(q_proj): Linear(in_features=16, out_features=64, bias=False)
(k_proj): Linear(in_features=16, out_features=64, bias=False)
(o_proj): Linear(in_features=64, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
)
(norm): DiffusionGemmaRMSNorm()
(rotary_emb): DiffusionGemmaTextRotaryEmbedding()
)
(vision_tower): Gemma4VisionModel(
(patch_embedder): Gemma4VisionPatchEmbedder(
(input_proj): Linear(in_features=768, out_features=64, bias=False)
)
(encoder): Gemma4VisionEncoder(
(rotary_emb): Gemma4VisionRotaryEmbedding()
(layers): ModuleList(
(0-1): 2 x Gemma4VisionEncoderLayer(
(self_attn): Gemma4VisionAttention(
(q_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=64, bias=False)
)
(k_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=64, bias=False)
)
(v_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=64, bias=False)
)
(o_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=64, bias=False)
)
(q_norm): Gemma4RMSNorm()
(k_norm): Gemma4RMSNorm()
(v_norm): Gemma4RMSNorm()
)
(mlp): Gemma4VisionMLP(
(gate_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=128, bias=False)
)
(up_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=64, out_features=128, bias=False)
)
(down_proj): Gemma4ClippableLinear(
(linear): Linear(in_features=128, out_features=64, bias=False)
)
(act_fn): GELUTanh()
)
(input_layernorm): Gemma4RMSNorm()
(post_attention_layernorm): Gemma4RMSNorm()
(pre_feedforward_layernorm): Gemma4RMSNorm()
(post_feedforward_layernorm): Gemma4RMSNorm()
)
)
)
(pooler): Gemma4VisionPooler()
)
(embed_vision): DiffusionGemmaMultimodalEmbedder(
(embedding_projection): Linear(in_features=64, out_features=16, bias=False)
(embedding_pre_projection_norm): DiffusionGemmaRMSNorm()
)
)
(decoder): DiffusionGemmaDecoderModel(
(embed_tokens): DiffusionGemmaTextScaledWordEmbedding(262144, 16, padding_idx=0)
(layers): ModuleList(
(0-4): 5 x DiffusionGemmaDecoderTextLayer(
(self_attn): DiffusionGemmaDecoderTextAttention(
(q_proj): Linear(in_features=16, out_features=32, bias=False)
(k_proj): Linear(in_features=16, out_features=32, bias=False)
(v_proj): Linear(in_features=16, out_features=32, bias=False)
(o_proj): Linear(in_features=32, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(5): DiffusionGemmaDecoderTextLayer(
(self_attn): DiffusionGemmaDecoderTextAttention(
(q_proj): Linear(in_features=16, out_features=64, bias=False)
(k_proj): Linear(in_features=16, out_features=64, bias=False)
(o_proj): Linear(in_features=64, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(6-10): 5 x DiffusionGemmaDecoderTextLayer(
(self_attn): DiffusionGemmaDecoderTextAttention(
(q_proj): Linear(in_features=16, out_features=32, bias=False)
(k_proj): Linear(in_features=16, out_features=32, bias=False)
(v_proj): Linear(in_features=16, out_features=32, bias=False)
(o_proj): Linear(in_features=32, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
(11): DiffusionGemmaDecoderTextLayer(
(self_attn): DiffusionGemmaDecoderTextAttention(
(q_proj): Linear(in_features=16, out_features=64, bias=False)
(k_proj): Linear(in_features=16, out_features=64, bias=False)
(o_proj): Linear(in_features=64, out_features=16, bias=False)
(q_norm): DiffusionGemmaRMSNorm()
(k_norm): DiffusionGemmaRMSNorm()
(v_norm): DiffusionGemmaRMSNorm()
)
(mlp): DiffusionGemmaText4MLP(
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
(input_layernorm): DiffusionGemmaRMSNorm()
(post_attention_layernorm): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm): DiffusionGemmaRMSNorm()
(router): DiffusionGemmaTextRouter(
(norm): DiffusionGemmaRMSNorm()
(proj): Linear(in_features=16, out_features=32, bias=False)
)
(experts): DiffusionGemmaTextExperts(
(act_fn): GELUTanh()
)
(post_feedforward_layernorm_1): DiffusionGemmaRMSNorm()
(post_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
(pre_feedforward_layernorm_2): DiffusionGemmaRMSNorm()
)
)
(norm): DiffusionGemmaRMSNorm()
(rotary_emb): DiffusionGemmaTextRotaryEmbedding()
(self_conditioning): DiffusionGemmaSelfConditioning(
(pre_norm): DiffusionGemmaRMSNorm()
(post_norm): DiffusionGemmaRMSNorm()
(gate_proj): Linear(in_features=16, out_features=64, bias=False)
(up_proj): Linear(in_features=16, out_features=64, bias=False)
(down_proj): Linear(in_features=64, out_features=16, bias=False)
(act_fn): GELUTanh()
)
)
)
(lm_head): Linear(in_features=16, out_features=262144, bias=False)
)
```