gemma/diffusion/hackable_diffusion_adapter/sampling_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "A100 80GB",
      "private_outputs": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Colab for DiffusionGemma sampling using HD\n",
        "\n",
        "Please note that the codebase powering this demo prioritizes clear, modifiable code to encourage experimentation and development. As a result, it has not been optimized for inference speed."
      ],
      "metadata": {},
      "id": "sRadOiQNvp-N"
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q gemma\n",
        "\n",
        "## This colab requires tensorflow 2.21.0:\n",
        "# !pip uninstall -y tensorflow-cpu\n",
        "!pip install tensorflow==2.21.0"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": [],
      "id": "H7qw8woo8GOv"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Imports\n",
        "\n",
        "# for checkpoint\n",
        "# Common imports\n",
        "import os\n",
        "from gemma import gm\n",
        "from gemma import diffusion\n",
        "\n",
        "# for sampling\n",
        "import time\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import dialog\n",
        "import numpy as np\n",
        "\n",
        "from gemma.gm import text as gm_text\n",
        "from hackable_diffusion import hd\n",
        "\n",
        "from gemma.diffusion import _models as gemma_diffusion\n",
        "from gemma.diffusion.hackable_diffusion_adapter.hd import hd_gemma_network\n",
        "from gemma.diffusion.hackable_diffusion_adapter.hd import hd_gemma_ar_state_handler"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": [],
      "id": "aF3iFUGvlUdp"
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize the Diffusion Gemma model\n",
        "gemma_model = diffusion.DiffusionGemma_26B_A4B()\n",
        "\n",
        "# Load pretrained parameters\n",
        "gemma_params = gm.ckpts.load_params(diffusion.CheckpointPath.DIFFUSIONGEMMA_26B_A4B_IT, restore_concurrent_gb=16)\n",
        "print(\"Checkpoint loaded.\")"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": [],
      "id": "rX67N_wejPhE"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## HD Sampler"
      ],
      "metadata": {},
      "id": "N1odQ5jkGtal"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Padding utility functions\n",
        "\n",
        "def format_prompt(\n",
        "    prompt: str,\n",
        "    format_str: dialog.Format,\n",
        "    system_instruction: str | None = None,\n",
        "    add_think_token: bool = False,\n",
        ") -> str:\n",
        "  \"\"\"Formats a prompt for Gemma.\"\"\"\n",
        "  transformations = []\n",
        "  system_transforms = []\n",
        "  if add_think_token:\n",
        "    system_transforms.append(dialog.Think())\n",
        "  if system_instruction is not None:\n",
        "    system_transforms.append(system_instruction)\n",
        "  if system_transforms:\n",
        "    transformations.append(dialog.System(*system_transforms))\n",
        "  transformations.append(dialog.User(prompt))\n",
        "  conv = dialog.Conversation(*transformations)\n",
        "  # See third_party/py/gemma/gm/text/_chat_sampler.py;l=322\n",
        "  return conv.as_text(format=format_str, add_tool_response_tag_after_call=True)\n",
        "\n",
        "\n",
        "def tokenize_prompts_right_pad_fixed_max_length(\n",
        "    tokenizer,\n",
        "    prompts,\n",
        "    system_instruction: str | None = None,\n",
        "    add_think_token: bool = False,\n",
        "    pad_token: int = 0,\n",
        "    max_length: int = 64,\n",
        ") -> tuple[jnp.ndarray, jnp.ndarray]:\n",
        "  \"\"\"Tokenizes and right-pads a batch of prompts to equal length.\n",
        "\n",
        "  Matches Gemma convention: real tokens first, padding at the end.\n",
        "\n",
        "  Args:\n",
        "    tokenizer: Tokenizer with a ``FORMAT`` attribute and ``encode`` method.\n",
        "    prompts: Batch of text prompts.\n",
        "    system_instruction: System instruction to prepend to the prompt.\n",
        "    add_think_token: Whether to add a think token to the system instruction.\n",
        "    pad_token: Token id used for padding.\n",
        "\n",
        "  Returns:\n",
        "    A tuple of:\n",
        "      - Padded token array of shape ``[B, max_prompt_len]`` (right-padded).\n",
        "      - Array of per-element prompt lengths of shape ``[B]``.\n",
        "  \"\"\"\n",
        "  if not prompts:\n",
        "    raise ValueError(\"Must have at least one prompt.\")\n",
        "\n",
        "  all_token_ids = []\n",
        "  for prompt in prompts:\n",
        "    formatted = format_prompt(\n",
        "        prompt,\n",
        "        format_str=tokenizer.FORMAT,\n",
        "        system_instruction=system_instruction,\n",
        "        add_think_token=add_think_token,\n",
        "    )\n",
        "    prompt_ids = tokenizer.encode(formatted, add_bos=True)\n",
        "    all_token_ids.append(prompt_ids)\n",
        "\n",
        "  prompt_lengths = jnp.array(\n",
        "      [len(ids) for ids in all_token_ids], dtype=jnp.int32\n",
        "  )\n",
        "  # max_len = max(len(ids) for ids in all_token_ids)\n",
        "\n",
        "  # Right-pad: real tokens first, pad tokens at the end.\n",
        "  padded = []\n",
        "  for ids in all_token_ids:\n",
        "    if len(ids) < max_length:\n",
        "      padded.append(ids + [pad_token] * (max_length - len(ids)))\n",
        "    else:\n",
        "      raise ValueError(f\"Prompt length {len(ids)} exceeds max length {max_length}.\")\n",
        "\n",
        "  return jnp.array(padded, dtype=jnp.int32), prompt_lengths"
      ],
      "metadata": {
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": [],
      "id": "g5cZpsWaRuah"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Trajectory Sampling"
      ],
      "metadata": {},
      "id": "3ZbJ10U2eRG5"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title  Define canvas sampler\n",
        "\n",
        "# 1. Load Model & Tokenizer\n",
        "gemma_model = gemma_diffusion.DiffusionGemma_26B_A4B()\n",
        "\n",
        "gemma_tokenizer = gm_text.Gemma4Tokenizer()\n",
        "print(\"Tokenizer ready.\")\n",
        "\n",
        "# 2. Define Sampling Constants\n",
        "CANVAS_LENGTH = 512        #@param {type:\"integer\"}\n",
        "MAX_DENOISING_STEPS = 20   #@param {type:\"integer\"}\n",
        "MAX_PROMPT_LENGTH = 64     #@param {type:\"integer\"}\n",
        "# 3. Setup Hackable Diffusion Process & Stepper\n",
        "diffusion_process = hd.corruption.CategoricalProcess.uniform_process(\n",
        "    schedule=hd.corruption.LinearDiscreteSchedule(),\n",
        "    num_categories=gemma_tokenizer.vocab_size,\n",
        ")\n",
        "\n",
        "stop_tokens = (\n",
        "    gemma_tokenizer.special_tokens.EOS,\n",
        "    gemma_tokenizer.special_tokens.END_OF_TURN,\n",
        "    gemma_tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE,\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "# Define stepper\n",
        "stepper = hd.sampling.DiscreteDDIMStep(\n",
        "    corruption_process=diffusion_process,\n",
        "    temperature=0.7,\n",
        "    logits_dtype=jnp.bfloat16,\n",
        ")\n",
        "\n",
        "canvas_sampler = hd.sampling.DiffusionSampler(\n",
        "    time_schedule=hd.sampling.UniformTimeSchedule(),\n",
        "    stepper=stepper,\n",
        "    num_steps=MAX_DENOISING_STEPS,\n",
        "    update_conditioning_fn=hd_gemma_ar_state_handler.PropagateSelfConditioningFn(),\n",
        "    store_trajectory=True\n",
        ")\n",
        "\n",
        "# 5. Define JITted Sampling Function\n",
        "# @jax.jit\n",
        "def sample_fn(params, rng_key, p_tokens, p_lengths):\n",
        "  \"\"\"JIT compiled autoregressive sampling function.\"\"\"\n",
        "  # Move instantiation inside so JAX doesn't try to hash them at the boundary\n",
        "  diffusion_network = hd_gemma_network.WrappedDiffusionGemmaNetwork(\n",
        "      gemma_model=gemma_model\n",
        "  )\n",
        "  inference_fn = hd.inference.FlaxLinenInferenceFn(\n",
        "      network=diffusion_network,\n",
        "      params={'gemma_model': params},\n",
        "  )\n",
        "  gemma_state_handler = hd_gemma_ar_state_handler.GemmaARStateHandler(\n",
        "    gemma_network=diffusion_network,\n",
        "    gemma_params=params,\n",
        "    end_tokens=stop_tokens,\n",
        "  )\n",
        "\n",
        "  sampler_state = gemma_state_handler.init_ar_state(\n",
        "      batch_size=len(p_tokens),\n",
        "      conditioning={\n",
        "          \"prompt_tokens\": p_tokens,\n",
        "          \"prompt_lengths\": p_lengths,\n",
        "      },\n",
        "      canvas_length=CANVAS_LENGTH,\n",
        "      max_num_canvases=1,\n",
        "  )\n",
        "  # Propagate random number generator.\n",
        "  canvas_init_rng, canvas_sampler_rng = jax.random.split(rng_key)\n",
        "\n",
        "  # Create new canvas.\n",
        "  initial_canvas = diffusion_process.sample_from_invariant(\n",
        "    key=canvas_init_rng,\n",
        "    data_spec=jnp.zeros((len(p_tokens), CANVAS_LENGTH, 1), dtype=jnp.int32),\n",
        "  )\n",
        "  # Sample canvas via diffusion.\n",
        "  return canvas_sampler(\n",
        "    inference_fn=inference_fn,\n",
        "    rng=canvas_sampler_rng,\n",
        "    initial_noise=initial_canvas,\n",
        "    conditioning=gemma_state_handler.create_conditioning_from_state(\n",
        "        sampler_state=sampler_state\n",
        "    ),\n",
        "  )\n",
        "\n",
        "class HDDiffusionSampler:\n",
        "  def __init__(self, params, tokenizer, max_length):\n",
        "    self.params = params\n",
        "    self.tokenizer = tokenizer\n",
        "    self.max_length = max_length\n",
        "\n",
        "  def sample(self, prompt, *, rng=None):\n",
        "    if isinstance(prompt, str):\n",
        "      prompt = [prompt]\n",
        "    if rng is None:\n",
        "      rng = jax.random.PRNGKey(42)\n",
        "\n",
        "    # Tokenize\n",
        "    prompt_tokens, prompt_lengths = tokenize_prompts_right_pad_fixed_max_length(\n",
        "        self.tokenizer, prompt\n",
        "    )\n",
        "\n",
        "    # Cast to int32 before padding to ensure JAX cache hit\n",
        "    prompt_tokens = jnp.array(prompt_tokens, dtype=jnp.int32)\n",
        "    prompt_lengths = jnp.array(prompt_lengths, dtype=jnp.int32)\n",
        "\n",
        "    # Execute\n",
        "    return sample_fn(self.params, rng, prompt_tokens, prompt_lengths)\n",
        "\n",
        "sampler = HDDiffusionSampler(\n",
        "    params=gemma_params,\n",
        "    tokenizer=gemma_tokenizer,\n",
        "    max_length=MAX_PROMPT_LENGTH\n",
        ")\n",
        "\n",
        "PROMPT = \"Summarize why text diffusion models are better than AR models\" # @param {type:\"string\"}\n",
        "rng = jax.random.PRNGKey(42)\n",
        "print(f\"User: {PROMPT}\")\n",
        "print(\"\\nSampling (This might take 2-3 minutes)...\")\n",
        "tic = time.time()\n",
        "last_step, trajectory = sampler.sample(PROMPT, rng=rng)\n",
        "print(f\"\\nSampling takes: {time.time()-tic:.1f}s.\")\n",
        "# === Extract results to CPU immediately ===\n",
        "final_tokens = np.array(last_step.xt[..., 0])           # numpy, on CPU\n",
        "trajectory_np = np.array(trajectory.xt[:, :, :, 0])  # (steps, batch, seq_len)\n",
        "trajectory_tokens = [trajectory_np[i, 0, :].tolist() for i in range(trajectory_np.shape[0])] # numpy, on CPU\n",
        "text = gemma_tokenizer.decode(final_tokens[0].tolist())\n",
        "print('\\nText Output:\\n-------------------------------------\\n', text)\n",
        "# === Free GPU memory ===\n",
        "del last_step, trajectory"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": [],
      "id": "w4kTn80Hv152"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Render GIF for trajectory visualization\n",
        "\n",
        "from PIL import Image as PILImage, ImageDraw, ImageFont\n",
        "import numpy as np\n",
        "import random\n",
        "\n",
        "GIBBERISH = list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%&*?+=<>~^')\n",
        "\n",
        "def make_diffusion_gif(\n",
        "    trajectory_tokens,\n",
        "    tokenizer,\n",
        "    output_path='diffusion_trajectory_scan.gif',\n",
        "    prompt=PROMPT,\n",
        "    font_size=28,\n",
        "    hold_ms=400,\n",
        "    chars_per_line=80,\n",
        "    margin=30,\n",
        "    num_scan_frames=12,\n",
        "    scan_ms=40,\n",
        "    glow_radius=12,\n",
        "    band_width=80,\n",
        "    last_frame_ms=5000,\n",
        "):\n",
        "  font = ImageFont.truetype(\"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf\", font_size)\n",
        "  title_font = ImageFont.truetype(\"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf\", font_size + 8)\n",
        "\n",
        "  tmp = PILImage.new('RGB', (1, 1))\n",
        "  tmp_draw = ImageDraw.Draw(tmp)\n",
        "  cell_w = tmp_draw.textbbox((0, 0), 'M', font=font)[2] + 2\n",
        "  cell_h = int(tmp_draw.textbbox((0, 0), 'Mg', font=font)[3] * 1.5)\n",
        "\n",
        "  def safe_decode(tok_id, use_whitespace=False):\n",
        "      text = tokenizer.decode([int(tok_id)])\n",
        "      text = text.replace('\\t', '→')\n",
        "      out = []\n",
        "      for ch in text:\n",
        "          if ch == '\\n':\n",
        "              out.append('\\n')\n",
        "          elif 32 <= ord(ch) < 127:\n",
        "              out.append(ch)\n",
        "          else:\n",
        "              out.append(' ' if use_whitespace else random.choice(GIBBERISH))\n",
        "      return ''.join(out) or '·'\n",
        "\n",
        "  num_steps = len(trajectory_tokens)\n",
        "  num_tokens = len(trajectory_tokens[0])\n",
        "  final_tokens = trajectory_tokens[-1]\n",
        "\n",
        "  all_decoded = []\n",
        "  for s in range(num_steps):\n",
        "      step_decoded = []\n",
        "      for tok_idx in range(num_tokens):\n",
        "          matches_final = trajectory_tokens[s][tok_idx] == final_tokens[tok_idx]\n",
        "          step_decoded.append(safe_decode(\n",
        "              trajectory_tokens[s][tok_idx],\n",
        "              use_whitespace=matches_final,\n",
        "          ))\n",
        "      all_decoded.append(step_decoded)\n",
        "  print(\"Tokens decoded.\")\n",
        "\n",
        "  # Final text defines the board\n",
        "  final_decoded = all_decoded[-1]\n",
        "  cells_per_token = [len(final_decoded[tok_idx]) for tok_idx in range(num_tokens)]\n",
        "  total_chars = sum(cells_per_token)\n",
        "\n",
        "  def fit_to_board(step_idx):\n",
        "      result = []\n",
        "      for tok_idx in range(num_tokens):\n",
        "          allocated = cells_per_token[tok_idx]\n",
        "          decoded = all_decoded[step_idx][tok_idx]\n",
        "          if len(decoded) <= allocated:\n",
        "              padded = decoded.ljust(allocated)\n",
        "          else:\n",
        "              padded = decoded[:allocated]\n",
        "          for ch in padded:\n",
        "              result.append((ch, tok_idx))\n",
        "      return result\n",
        "\n",
        "  GREEN = (80, 220, 100)\n",
        "  RED = (255, 80, 80)\n",
        "  GREY = (130, 130, 140)\n",
        "  BG = (25, 25, 35)\n",
        "  WHITE = (220, 220, 235)\n",
        "  GLOW_COLOR = np.array([80, 180, 255], dtype=np.float32)\n",
        "\n",
        "  # Pre-compute all boards\n",
        "  all_boards = [fit_to_board(s) for s in range(num_steps)]\n",
        "\n",
        "  # Compute grid positions respecting \\n in final text\n",
        "  final_board = all_boards[-1]\n",
        "  positions = []\n",
        "  col = 0\n",
        "  row = 0\n",
        "  for i, (ch, tok_idx) in enumerate(final_board):\n",
        "      if ch == '\\n':\n",
        "          positions.append(None)\n",
        "          col = 0\n",
        "          row += 1\n",
        "      else:\n",
        "          positions.append((col, row))\n",
        "          col += 1\n",
        "          if col >= chars_per_line:\n",
        "              col = 0\n",
        "              row += 1\n",
        "\n",
        "  actual_num_lines = row + 1\n",
        "  header_h = cell_h * (4 if prompt else 3)\n",
        "  img_w = margin * 2 + chars_per_line * cell_w\n",
        "  img_h = header_h + actual_num_lines * cell_h + margin\n",
        "\n",
        "  print(f\"Board: {total_chars} cells, {actual_num_lines} lines\")\n",
        "\n",
        "  def compute_colors(step_idx, prev_tokens, prev_colors):\n",
        "      cur_tokens = trajectory_tokens[step_idx]\n",
        "      colors = []\n",
        "      for tok_idx in range(num_tokens):\n",
        "          changed = prev_tokens is not None and cur_tokens[tok_idx] != prev_tokens[tok_idx]\n",
        "          matches_final = cur_tokens[tok_idx] == final_tokens[tok_idx]\n",
        "          if matches_final:\n",
        "              colors.append('green')\n",
        "          elif changed:\n",
        "              colors.append('red')\n",
        "          elif prev_colors and prev_colors[tok_idx] == 'green':\n",
        "              colors.append('green')\n",
        "          else:\n",
        "              colors.append('grey')\n",
        "      return colors\n",
        "\n",
        "  def render_full_frame(board, tok_colors, step_idx):\n",
        "      img = PILImage.new('RGB', (img_w, img_h), color=BG)\n",
        "      draw = ImageDraw.Draw(img)\n",
        "\n",
        "      if prompt:\n",
        "          draw.text((margin, margin), f'Prompt: {prompt}',\n",
        "                    fill=(180, 180, 200), font=font)\n",
        "\n",
        "      title_y = margin + cell_h if prompt else margin\n",
        "      draw.text((margin, title_y), f'Denoising Step {step_idx}/{num_steps - 1}',\n",
        "                fill=WHITE, font=title_font)\n",
        "\n",
        "      n_green = sum(1 for c in tok_colors if c == 'green')\n",
        "      pct = int(n_green / num_tokens * 100)\n",
        "      draw.text((margin + 500, title_y), f'{pct}% converged',\n",
        "                fill=(100, 100, 120), font=font)\n",
        "\n",
        "      bar_y = title_y + cell_h + 10\n",
        "      bar_w = img_w - margin * 2\n",
        "      draw.rounded_rectangle([margin, bar_y, margin + bar_w, bar_y + 10],\n",
        "                             radius=5, fill=(60, 60, 70))\n",
        "      progress = step_idx / max(num_steps - 1, 1)\n",
        "      if progress > 0:\n",
        "          draw.rounded_rectangle([margin, bar_y, margin + int(bar_w * progress), bar_y + 10],\n",
        "                                 radius=5, fill=(100, 180, 255))\n",
        "\n",
        "      color_map = {'green': GREEN, 'red': RED, 'grey': GREY}\n",
        "      for i, (ch, tok_idx) in enumerate(board):\n",
        "          if positions[i] is None:\n",
        "              continue\n",
        "          col_idx, row_idx = positions[i]\n",
        "          x = margin + col_idx * cell_w\n",
        "          y = header_h + row_idx * cell_h\n",
        "          color = color_map[tok_colors[tok_idx]] if 0 <= tok_idx < len(tok_colors) else GREY\n",
        "          if ch.strip():\n",
        "              cw = draw.textbbox((0, 0), ch, font=font)[2]\n",
        "              x_offset = (cell_w - cw) // 2\n",
        "              draw.text((x + x_offset, y), ch, fill=color, font=font)\n",
        "\n",
        "      return img\n",
        "\n",
        "  # Pre-compute vertical gradient lookup\n",
        "  ys = np.arange(img_h, dtype=np.float32)\n",
        "\n",
        "  def make_scanline_frame(old_arr, new_arr, scan_y):\n",
        "      alpha = np.clip((scan_y - ys) / band_width + 0.5, 0.0, 1.0)\n",
        "      alpha = alpha * alpha * (3 - 2 * alpha)\n",
        "      alpha = alpha.reshape(-1, 1, 1)\n",
        "\n",
        "      blended = new_arr * alpha + old_arr * (1 - alpha)\n",
        "\n",
        "      for offset in range(-2, 3):\n",
        "          y = scan_y + offset\n",
        "          if 0 <= y < img_h:\n",
        "              intensity = 1.0 - abs(offset) / 3\n",
        "              blended[y, :] = np.clip(\n",
        "                  blended[y, :] + GLOW_COLOR * intensity * 0.4, 0, 255\n",
        "              )\n",
        "\n",
        "      return PILImage.fromarray(blended.astype(np.uint8))\n",
        "\n",
        "  # Compute all colors\n",
        "  all_colors = []\n",
        "  prev_tokens = None\n",
        "  prev_colors = None\n",
        "  for s in range(num_steps):\n",
        "      colors = compute_colors(s, prev_tokens, prev_colors)\n",
        "      all_colors.append(colors)\n",
        "      prev_tokens = trajectory_tokens[s]\n",
        "      prev_colors = colors\n",
        "\n",
        "  # Render static frames\n",
        "  step_frames = []\n",
        "  for step_idx in range(num_steps):\n",
        "      print(f'\\rRendering frame {step_idx+1}/{num_steps}', end='')\n",
        "      step_frames.append(render_full_frame(all_boards[step_idx], all_colors[step_idx], step_idx))\n",
        "  print(' Done!')\n",
        "\n",
        "  # Convert to numpy once\n",
        "  step_arrays = [np.array(f, dtype=np.float32) for f in step_frames]\n",
        "\n",
        "  # Build animation with scanline transitions\n",
        "  frames = []\n",
        "  durations = []\n",
        "  print(\"Building transitions...\")\n",
        "\n",
        "  for step_idx in range(num_steps):\n",
        "      print(f'\\rStep {step_idx+1}/{num_steps}', end='')\n",
        "\n",
        "      if step_idx > 0:\n",
        "          has_changes = any(\n",
        "              trajectory_tokens[step_idx][i] != trajectory_tokens[step_idx - 1][i]\n",
        "              or all_colors[step_idx][i] != all_colors[step_idx - 1][i]\n",
        "              for i in range(num_tokens)\n",
        "          )\n",
        "\n",
        "          if has_changes:\n",
        "              old_arr = step_arrays[step_idx - 1]\n",
        "              new_arr = step_arrays[step_idx]\n",
        "              for f in range(num_scan_frames):\n",
        "                  scan_y = int(header_h + (img_h - header_h) * f / num_scan_frames)\n",
        "                  frames.append(make_scanline_frame(old_arr, new_arr, scan_y))\n",
        "                  durations.append(scan_ms)\n",
        "\n",
        "      frames.append(step_frames[step_idx])\n",
        "      durations.append(hold_ms)\n",
        "\n",
        "  durations[-1] = last_frame_ms\n",
        "  print(' Done!')\n",
        "\n",
        "  frames[0].save(output_path, save_all=True, append_images=frames[1:],\n",
        "                 duration=durations, loop=0)\n",
        "  print(f\"Saved to {output_path} ({len(frames)} frames from {num_steps} steps)\")\n",
        "\n",
        "\n",
        "make_diffusion_gif(trajectory_tokens, gemma_tokenizer, prompt=PROMPT)\n",
        "\n",
        "from IPython.display import display, Image\n",
        "display(Image(filename='diffusion_trajectory_scan.gif'))"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": [],
      "id": "N1BYJbts-hdN"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Full AR Sampling"
      ],
      "metadata": {},
      "id": "-O1iwzmfZwNG"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Define HD Sampler and sampling constants\n",
        "\n",
        "# 1. Load Model & Tokenizer\n",
        "gemma_model = gemma_diffusion.DiffusionGemma_26B_A4B()\n",
        "\n",
        "gemma_tokenizer = gm_text.Gemma4Tokenizer()\n",
        "print(\"Tokenizer ready.\")\n",
        "\n",
        "# 2. Define Sampling Constants\n",
        "CANVAS_LENGTH = 128        #@param {type:\"integer\"}\n",
        "MAX_DENOISING_STEPS = 20   #@param {type:\"integer\"}\n",
        "MAX_NUM_CANVASES = 8       #@param {type:\"integer\"}\n",
        "MAX_PROMPT_LENGTH = 64     #@param {type:\"integer\"}\n",
        "\n",
        "# 3. Decide the components\n",
        "stepper_type = \"DiscreteDDIM\" #@param [\"DiscreteDDIM\", \"IntegratedDiscreteDDIM\", \"DiscreteFlowMatching\"]\n",
        "routing_strategy = \"None\" # @param [\"None\", \"Greedy\"]\n",
        "use_early_stopping = True # @param {type:\"boolean\"}\n",
        "\n",
        "# 3. Setup Hackable Diffusion Process & Stepper\n",
        "diffusion_process = hd.corruption.CategoricalProcess.uniform_process(\n",
        "    schedule=hd.corruption.LinearDiscreteSchedule(),\n",
        "    num_categories=gemma_tokenizer.vocab_size,\n",
        ")\n",
        "\n",
        "stop_tokens = (\n",
        "    gemma_tokenizer.special_tokens.EOS,\n",
        "    gemma_tokenizer.special_tokens.END_OF_TURN,\n",
        "    gemma_tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE,\n",
        ")\n",
        "\n",
        "# Define early stopping strategy\n",
        "if use_early_stopping:\n",
        "  early_stop_fn = hd.sampling.DiffusionEntropyEarlyStopFn()\n",
        "else:\n",
        "  early_stop_fn = hd.sampling.DiffusionNoEarlyStopFn()\n",
        "\n",
        "\n",
        "# Define routing strategy\n",
        "match routing_strategy:\n",
        "  case \"None\":\n",
        "    planner = None\n",
        "  case \"Greedy\":\n",
        "    planner = hd.sampling.GreedyPlanner()\n",
        "\n",
        "\n",
        "# Define stepper\n",
        "match stepper_type:\n",
        "  case \"DiscreteDDIM\":\n",
        "    stepper = hd.sampling.DiscreteDDIMStep(\n",
        "        corruption_process=diffusion_process,\n",
        "        planner=planner,\n",
        "        temperature=0.7,\n",
        "        logits_dtype=jnp.bfloat16,\n",
        "    )\n",
        "  case \"IntegratedDiscreteDDIM\":\n",
        "    stepper = hd.sampling.IntegratedDiscreteDDIMStep(\n",
        "        corruption_process=diffusion_process,\n",
        "        temperature=0.7,\n",
        "        logits_dtype=jnp.bfloat16,\n",
        "    )\n",
        "  case \"DiscreteFlowMatching\":\n",
        "    stepper = hd.sampling.DiscreteFlowMatchingStep(\n",
        "        corruption_process=diffusion_process,\n",
        "        planner=planner,\n",
        "        temperature=0.7,\n",
        "        logits_dtype=jnp.bfloat16,\n",
        "    )\n",
        "\n",
        "\n",
        "# 4. Create Canvas Sampler\n",
        "canvas_sampler = hd.sampling.DiffusionSamplerWithEarlyStopping(\n",
        "    time_schedule=hd.sampling.UniformTimeSchedule(),\n",
        "    stepper=stepper,\n",
        "    num_steps=MAX_DENOISING_STEPS,\n",
        "    update_conditioning_fn=hd_gemma_ar_state_handler.PropagateSelfConditioningFn(),\n",
        "    early_stopping_fn=hd.sampling.DiffusionEntropyEarlyStopFn(),\n",
        ")\n",
        "\n",
        "# 5. Define JITted Sampling Function\n",
        "@jax.jit\n",
        "def sample_ar_fn(params, rng_key, p_tokens, p_lengths):\n",
        "  \"\"\"JIT compiled autoregressive sampling function.\"\"\"\n",
        "  # Move instantiation inside so JAX doesn't try to hash them at the boundary\n",
        "  diffusion_network = hd_gemma_network.WrappedDiffusionGemmaNetwork(\n",
        "      gemma_model=gemma_model\n",
        "  )\n",
        "  inference_fn = hd.inference.FlaxLinenInferenceFn(\n",
        "      network=diffusion_network,\n",
        "      params={'gemma_model': params},\n",
        "  )\n",
        "  gemma_state_handler = hd_gemma_ar_state_handler.GemmaARStateHandler(\n",
        "    gemma_network=diffusion_network,\n",
        "    gemma_params=params,\n",
        "    end_tokens=stop_tokens,\n",
        "  )\n",
        "  autoregressive_sampler = hd.sampling.AutoregressiveDiffusionSampler(\n",
        "      diffusion_process=diffusion_process,\n",
        "      canvas_sampler=canvas_sampler,\n",
        "      max_num_canvases=MAX_NUM_CANVASES,\n",
        "      canvas_length=CANVAS_LENGTH,\n",
        "      state_handler=gemma_state_handler,\n",
        "      data_dtype=jnp.int32,\n",
        "      data_shape=(1,),\n",
        "  )\n",
        "\n",
        "  return autoregressive_sampler(\n",
        "      diffusion_inference_fn=inference_fn,\n",
        "      batch_size=len(p_tokens),\n",
        "      rng=rng_key,\n",
        "      conditioning={\n",
        "          \"prompt_tokens\": p_tokens,\n",
        "          \"prompt_lengths\": p_lengths,\n",
        "      },\n",
        "  )\n",
        "\n",
        "class HDDiffusionSampler:\n",
        "  def __init__(self, params, tokenizer, max_length):\n",
        "    self.params = params\n",
        "    self.tokenizer = tokenizer\n",
        "    self.max_length = max_length\n",
        "\n",
        "  def sample(self, prompt, *, rng=None):\n",
        "    if isinstance(prompt, str):\n",
        "      prompt = [prompt]\n",
        "    if rng is None:\n",
        "      rng = jax.random.PRNGKey(42)\n",
        "\n",
        "    # Tokenize\n",
        "    prompt_tokens, prompt_lengths = tokenize_prompts_right_pad_fixed_max_length(\n",
        "        self.tokenizer, prompt\n",
        "    )\n",
        "\n",
        "    # Cast to int32 before padding to ensure JAX cache hit\n",
        "    prompt_tokens = jnp.array(prompt_tokens, dtype=jnp.int32)\n",
        "    prompt_lengths = jnp.array(prompt_lengths, dtype=jnp.int32)\n",
        "\n",
        "    # Execute\n",
        "    return sample_ar_fn(self.params, rng, prompt_tokens, prompt_lengths)\n",
        "\n",
        "sampler = HDDiffusionSampler(\n",
        "    params=gemma_params,\n",
        "    tokenizer=gemma_tokenizer,\n",
        "    max_length=MAX_PROMPT_LENGTH\n",
        ")\n",
        "\n",
        "rng = jax.random.PRNGKey(42)\n",
        "tic = time.time()\n",
        "ar_diffusion_samples = sampler.sample(\"Hello\", rng=rng)\n",
        "print(f\"\\nSampling takes: {time.time()-tic:.1f}s.\")\n",
        "out_text_ar = gemma_tokenizer.decode(ar_diffusion_samples[0][0])\n",
        "del ar_diffusion_samples\n",
        "print('\\nText Output:\\n-------------------------------------\\n', out_text_ar)\n",
        "print(\"HD Sampler ready.\")"
      ],
      "metadata": {
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": [],
      "id": "11FA4RpS2Bsg"
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "# @title Sampling #1\n",
        "PROMPT = \"What is the meaning of life and the rest?\" # @param {type:\"string\"}\n",
        "\n",
        "print(f\"\\nSampling: {PROMPT!r}\")\n",
        "rng = jax.random.PRNGKey(42)\n",
        "tic = time.time()\n",
        "ar_diffusion_samples = sampler.sample(PROMPT, rng=rng)\n",
        "print(f\"\\nSampling takes: {time.time()-tic:.1f}s.\")\n",
        "print(\"\\n--- Autoregressive Multi-Canvas Sampling Result ---\")\n",
        "out_text_ar = gemma_tokenizer.decode(ar_diffusion_samples[0][0])\n",
        "del ar_diffusion_samples\n",
        "print('\\nText Output:\\n-------------------------------------\\n', out_text_ar)"
      ],
      "metadata": {
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": [],
      "id": "z3lA1xIumi9W"
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Sampling #2\n",
        "PROMPT = \"What are the advantages of text diffusion models over AR models?\" # @param {type:\"string\"}\n",
        "\n",
        "print(f\"\\nSampling: {PROMPT!r}\")\n",
        "rng, _ = jax.random.split(rng)\n",
        "tic = time.time()\n",
        "ar_diffusion_samples = sampler.sample(PROMPT, rng=rng)\n",
        "print(f\"\\nSampling takes: {time.time()-tic:.1f}s.\")\n",
        "print(\"\\n--- Autoregressive Multi-Canvas Sampling Result ---\")\n",
        "out_text_ar = gemma_tokenizer.decode(ar_diffusion_samples[0][0])\n",
        "del ar_diffusion_samples\n",
        "print('\\nText Output:\\n-------------------------------------\\n', out_text_ar)"
      ],
      "metadata": {
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": [],
      "id": "DpHARBBH6pkm"
    }
  ]
}