Upload folder using huggingface_hub

README.md

@@ -1,13 +1,16 @@
 ---
-title: Diffusiongemma 26B A4B It
-emoji: 📚
-colorFrom: gray
-colorTo: red
+title: DiffusionGemma 26B A4B
+emoji: 🌀
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.18.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Block-diffusion chat with live denoising canvas
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DiffusionGemma 26B-A4B
+
+Chat demo for [google/diffusiongemma-26B-A4B-it](https://huggingface.co/google/diffusiongemma-26B-A4B-it) with a real-time visualization of the block-diffusion denoising canvas. Supports image input and a thinking mode.

app.py

@@ -0,0 +1,224 @@
+import spaces
+import torch
+import gradio as gr
+from threading import Thread
+from queue import Queue
+from transformers import DiffusionGemmaForBlockDiffusion, AutoProcessor, TextDiffusionStreamer
+
+MODEL_ID = "google/diffusiongemma-26B-A4B-it"
+
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = DiffusionGemmaForBlockDiffusion.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
+model.to("cuda")
+model.eval()
+
+_SENTINEL = object()
+
+
+class CanvasStreamer(TextDiffusionStreamer):
+    """Pushes (committed_text, draft_text) snapshots to a queue for live visualization.
+
+    `put_draft` fires every denoising step with the full canvas of the block being
+    denoised (the text morphs as it denoises). `put` fires when a canvas/block is
+    committed. The prompt is skipped via `skip_prompt`.
+    """
+
+    def __init__(self, tokenizer, **kwargs):
+        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True, **kwargs)
+        self.queue = Queue()
+        self.committed = ""
+        self._takes_logits = False
+
+    def put_draft(self, value, **kwargs):
+        if len(value.shape) > 1:
+            value = value[0]
+        draft = self.tokenizer.decode(value, skip_special_tokens=True)
+        self.queue.put((self.committed, draft))
+
+    def put(self, value):
+        if len(value.shape) > 1 and value.shape[0] > 1:
+            raise ValueError("batch size 1 only")
+        elif len(value.shape) > 1:
+            value = value[0]
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+            return
+        self.committed += self.tokenizer.decode(value, skip_special_tokens=True)
+        self.queue.put((self.committed, ""))
+
+    def end(self):
+        self.next_tokens_are_prompt = True
+        self.queue.put(_SENTINEL)
+
+
+def build_highlight(committed, draft):
+    """Committed text green ('high'), in-progress draft canvas yellow ('mid')."""
+    out = []
+    if committed:
+        out.append((committed, "high"))
+    if draft:
+        out.append((draft, "mid"))
+    return out
+
+
+@spaces.GPU(duration=150, size="xlarge")
+@torch.no_grad()
+def generate_streaming(messages, max_new_tokens, max_denoising_steps, enable_thinking):
+    inputs = processor.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+        enable_thinking=enable_thinking,
+    ).to("cuda")
+    prompt_len = inputs["input_ids"].shape[1]
+
+    streamer = CanvasStreamer(processor.tokenizer)
+    result = {}
+
+    def run():
+        try:
+            out = model.generate(
+                **inputs,
+                max_new_tokens=int(max_new_tokens),
+                max_denoising_steps=int(max_denoising_steps),
+                streamer=streamer,
+            )
+            result["ids"] = out
+        except Exception as e:
+            result["error"] = e
+            streamer.queue.put(_SENTINEL)
+
+    thread = Thread(target=run)
+    thread.start()
+
+    while True:
+        item = streamer.queue.get()
+        if item is _SENTINEL:
+            break
+        committed, draft = item
+        yield build_highlight(committed, draft), (committed + draft), None
+
+    thread.join()
+    if "error" in result:
+        raise result["error"]
+
+    generated = result["ids"][0][prompt_len:]
+    final_text = processor.decode(generated, skip_special_tokens=True).strip()
+    yield build_highlight(final_text, ""), final_text, final_text
+
+
+def to_model_messages(history):
+    """Convert Gradio messages history into processor chat format with images."""
+    messages = []
+    for msg in history:
+        role = msg["role"]
+        content = msg["content"]
+        if isinstance(content, str):
+            parts = [{"type": "text", "text": content}]
+        elif isinstance(content, (list, tuple)):
+            # Gradio file tuple (path, alt) or list of content dicts
+            parts = []
+            if isinstance(content, tuple):
+                parts.append({"type": "image", "url": content[0]})
+            else:
+                for item in content:
+                    if isinstance(item, dict) and item.get("type") == "text":
+                        parts.append({"type": "text", "text": item["text"]})
+                    elif isinstance(item, dict) and "path" in item:
+                        parts.append({"type": "image", "url": item["path"]})
+        else:
+            parts = [{"type": "text", "text": str(content)}]
+        messages.append({"role": role, "content": parts})
+    return messages
+
+
+css = """
+.category-legend{display:none}
+.legend{margin-bottom: 5px}
+.legend-item{height: 25px}
+"""
+
+
+def create_demo():
+    with gr.Blocks(css=css, title="DiffusionGemma") as demo:
+        gr.Markdown("# DiffusionGemma 26B-A4B — Block Diffusion Chat")
+        gr.Markdown(
+            "[model](https://huggingface.co/google/diffusiongemma-26B-A4B-it) · "
+            "Watch the canvas denoise in real time on the right. Attach an image to ask about it."
+        )
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(label="Conversation", type="messages", height=520)
+                chat_input = gr.MultimodalTextbox(
+                    interactive=True,
+                    file_types=["image"],
+                    placeholder="Type a message and/or attach an image…",
+                    show_label=False,
+                )
+            with gr.Column(scale=2):
+                canvas = gr.HighlightedText(
+                    label="Denoising canvas",
+                    combine_adjacent=False,
+                    show_legend=True,
+                    color_map={"mid": "#FFAA33", "high": "#66CC66"},
+                )
+
+        with gr.Accordion("Generation settings", open=False):
+            with gr.Row():
+                enable_thinking = gr.Checkbox(value=False, label="Thinking mode")
+            with gr.Row():
+                max_new_tokens = gr.Slider(64, 1024, value=256, step=64, label="Max new tokens")
+                max_denoising_steps = gr.Slider(8, 64, value=48, step=4, label="Max denoising steps")
+
+        clear_btn = gr.Button("Clear conversation")
+
+        def add_message(message, history):
+            history = history or []
+            for f in message.get("files", []):
+                history.append({"role": "user", "content": {"path": f}})
+            if message.get("text"):
+                history.append({"role": "user", "content": message["text"]})
+            return history, gr.MultimodalTextbox(value=None, interactive=False)
+
+        def bot(history, max_new_tokens, max_denoising_steps, enable_thinking):
+            if not history:
+                yield history, []
+                return
+            messages = to_model_messages(history)
+            history = history + [{"role": "assistant", "content": ""}]
+            final = None
+            try:
+                for canvas_state, plain, text in generate_streaming(
+                    messages, max_new_tokens, max_denoising_steps, enable_thinking
+                ):
+                    if text is not None:
+                        final = text
+                    history[-1]["content"] = final if final is not None else plain
+                    yield history, canvas_state
+            except Exception as e:
+                history[-1]["content"] = f"Error: {e}"
+                yield history, [(str(e), "mid")]
+
+        def reenable():
+            return gr.MultimodalTextbox(interactive=True)
+
+        chat_msg = chat_input.submit(
+            add_message, [chat_input, chatbot], [chatbot, chat_input]
+        )
+        bot_msg = chat_msg.then(
+            bot,
+            [chatbot, max_new_tokens, max_denoising_steps, enable_thinking],
+            [chatbot, canvas],
+        )
+        bot_msg.then(reenable, None, [chat_input])
+
+        clear_btn.click(lambda: ([], []), None, [chatbot, canvas])
+
+    return demo
+
+
+if __name__ == "__main__":
+    create_demo().queue().launch()

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+accelerate
+torchvision