How to use from
Docker Model Runner
docker model run hf.co/tiny-random/diffusiongemma-moe
Quick Links

This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from 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)
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
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)

Printing the model:

Click to expand
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)
)

Test environment:

  • Pillow: 13.0.0
  • huggingface_hub: 2.3.0
  • numpy: 2.3.5
  • safetensors: 0.8.0
  • tokenizers: 0.23.3
  • torch: 2.13.0
  • transformers: 5.19.0

Change log:

  • 2026-10-23: Initial version.

This repo is generated by GPT-6.1-Sol

Downloads last month
22
Safetensors
Model size
6.3M params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for tiny-random/diffusiongemma-moe

Finetuned
(37)
this model

Collection including tiny-random/diffusiongemma-moe