{ "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" } ] }