Upload folder using huggingface_hub

README.md

@@ -1,30 +1,31 @@
 ---
-title: Grug-9B Reasoning VLM
-emoji: 🧠
+title: Grug-9B Demo
+emoji: 🪨
 colorFrom: red
 colorTo: yellow
 sdk: gradio
-sdk_version: 6.20.0
+sdk_version: 6.15.1
 app_file: app.py
-short_description: 9B reasoning VLM with vision and shortened thinking
+short_description: Vision-language reasoning model with concise thinking
 python_version: "3.12"
 startup_duration_timeout: 1h
 ---
 
-# Grug-9B Reasoning VLM
+# 🪨 Grug-9B Vision-Language Demo
 
-A demo for [ProCreations/grug-9b](https://huggingface.co/ProCreations/grug-9b) — a 9B-parameter vision-language reasoning model fine-tuned from Ornith-1.0-9B (Qwen3.5 architecture) to produce shorter internal reasoning while maintaining coding and agent capabilities.
+A Gradio demo for [ProCreations/grug-9b](https://huggingface.co/ProCreations/grug-9b) —
+a 9B-parameter reasoning VLM fine-tuned from Ornith-1.0-9B (Qwen3.5 architecture)
+that produces concise reasoning traces ("grug think small") instead of verbose chain-of-thought.
 
 ## Features
 
-- **Vision + Language**: Upload an image and ask questions about it.
-- **Reasoning**: The model thinks inside `` tags before producing its answer.
-- **Chat**: Multi-turn conversation with the model.
-- **Code**: Ask the model to write functions, explain concepts, etc.
+- **Image + text input**: Upload an image and ask questions about it
+- **Thinking mode**: Toggle the model's reasoning traces (default on)
+- **Adjustable generation**: Max tokens, temperature, top-p sampling
 
-## Usage
+## Model details
 
-1. Optionally upload an image.
-2. Type your message in the text box.
-3. Press Send or hit Enter.
-4. Adjust advanced settings (max tokens, temperature, thinking mode) as needed.
+- Architecture: Qwen3.5 (`Qwen3_5ForConditionalGeneration`)
+- Parameters: ~9.4B (dense, bf16)
+- Base model: [deepreinforce-ai/Ornith-1.0-9B](https://huggingface.co/deepreinforce-ai/Ornith-1.0-9B)
+- License: MIT

app.py

@@ -2,128 +2,179 @@ import os
 
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-import spaces  # MUST come before torch / transformers
+import spaces  # noqa: E402  -- must be before torch
 import torch
 import gradio as gr
-import numpy as np
-from PIL import Image
-from threading import Thread
-from transformers import (
-    AutoProcessor,
-    Qwen3_5ForConditionalGeneration,
-    TextIteratorStreamer,
-)
+from transformers import AutoModelForImageTextToText, AutoProcessor
+
 
 MODEL_ID = "ProCreations/grug-9b"
+PROCESSOR_ID = "deepreinforce-ai/Ornith-1.0-9B"  # grug-9b omits preprocessor configs; base model has them
+
+print(f"Loading processor from {PROCESSOR_ID} …")
+processor = AutoProcessor.from_pretrained(PROCESSOR_ID)
 
-print(f"Loading model {MODEL_ID} ...")
-processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = Qwen3_5ForConditionalGeneration.from_pretrained(
+print(f"Loading model from {MODEL_ID} …")
+model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
-).to("cuda").eval()
+    dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+).to("cuda")
+model.eval()
 print("Model loaded.")
 
 
+def _build_messages(history, image_path, user_text):
+    """Build the messages list from chat history + new user input."""
+    messages = []
+    for msg in history:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        if isinstance(content, list):
+            messages.append({"role": role, "content": content})
+        else:
+            messages.append({"role": role, "content": [{"type": "text", "text": str(content)}]})
+    # Add the new user message
+    user_content = []
+    if image_path is not None:
+        user_content.append({"type": "image", "image": image_path})
+    user_content.append({"type": "text", "text": user_text})
+    messages.append({"role": "user", "content": user_content})
+    return messages
+
+
+@spaces.GPU(duration=120)
+def predict(image_path, user_text, max_new_tokens, temperature, top_p, enable_thinking):
+    """Run a single-turn vision+text inference and return the response.
+
+    Args:
+        image_path: path to the uploaded image (or None for text-only).
+        user_text: the user's text prompt.
+        max_new_tokens: maximum number of tokens to generate.
+        temperature: sampling temperature (1.0 = greedy-ish).
+        top_p: nucleus sampling probability.
+        enable_thinking: whether to emit <think> reasoning before the answer.
+
+    Returns:
+        The decoded text response.
+    """
+    if not user_text.strip() and image_path is None:
+        return "Please provide some text or an image to analyze."
+
+    messages = _build_messages([], image_path, user_text if user_text.strip() else "Describe this image.")
+
+    text = processor.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=enable_thinking,
+    )
+
+    # Build image inputs from the messages
+    image_inputs = []
+    for msg in messages:
+        for item in msg.get("content", []):
+            if isinstance(item, dict) and item.get("type") == "image":
+                from PIL import Image
+                image_inputs.append(Image.open(item["image"]).convert("RGB"))
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs if image_inputs else None,
+        padding=True,
+        return_tensors="pt",
+    ).to("cuda")
+
+    do_sample = temperature > 0.001
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=do_sample,
+            temperature=float(temperature) if do_sample else 1.0,
+            top_p=float(top_p) if do_sample else 1.0,
+        )
+
+    # Strip the input tokens from the output
+    generated = output_ids[0][inputs["input_ids"].shape[1]:]
+    result = processor.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    return result
+
+
 @spaces.GPU(duration=120)
-def chat_with_image(
-    image: Image.Image,
-    message: str,
-    history: list,
-    max_new_tokens: int = 1024,
-    temperature: float = 0.7,
-    enable_thinking: bool = True,
-):
-    """Chat with the grug-9b reasoning VLM about an image.
+def chat_predict(message, history, image, max_new_tokens, temperature, top_p, enable_thinking):
+    """Multi-turn chat with optional image. Gradio passes history as a list of
+    [user, assistant] tuples.
 
     Args:
-        image: An optional image to discuss with the model.
-        message: The user's text message / question.
-        history: Prior conversation turns (list of [user, assistant] pairs).
-        max_new_tokens: Maximum number of tokens to generate.
-        temperature: Sampling temperature; lower = more deterministic.
-        enable_thinking: Whether the model produces reasoning inside  tags.
+        message: the user's latest text message.
+        history: list of (user_msg, assistant_msg) tuples.
+        image: optional uploaded image path.
+        max_new_tokens: maximum tokens to generate.
+        temperature: sampling temperature.
+        top_p: nucleus sampling probability.
+        enable_thinking: whether to emit <think> reasoning.
+
+    Returns:
+        The assistant's text response.
     """
-    if not message or not message.strip():
-        yield history + [(message, "")], "Please enter a message."
-        return
+    if not message.strip() and image is None:
+        return "Please provide a message or an image."
 
-    # Build the conversation
+    # Convert Gradio history format to messages
     messages = []
+    for user_msg, assistant_msg in history:
+        # Rebuild each prior turn as a content list
+        user_content = []
+        # We can't perfectly reconstruct which prior messages had images,
+        # so we store images as text references in history
+        user_content.append({"type": "text", "text": user_msg})
+        messages.append({"role": "user", "content": user_content})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]})
+
+    # Add the current message
+    user_content = []
     if image is not None:
-        messages.append({
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": message},
-            ],
-        })
-    else:
-        messages.append({"role": "user", "content": message})
-
-    # Add conversation history
-    conv_messages = []
-    for user_msg, asst_msg in history:
-        conv_messages.append({"role": "user", "content": user_msg})
-        if asst_msg:
-            conv_messages.append({"role": "assistant", "content": asst_msg})
-    conv_messages.extend(messages)
-
-    # Prepare inputs
-    if image is not None:
-        text = processor.apply_chat_template(
-            conv_messages,
-            tokenize=False,
-            add_generation_prompt=True,
-            enable_thinking=enable_thinking,
-        )
-        inputs = processor(
-            text=[text],
-            images=[image],
-            return_tensors="pt",
-            padding=True,
-        ).to("cuda")
-    else:
-        text = processor.apply_chat_template(
-            conv_messages,
-            tokenize=False,
-            add_generation_prompt=True,
-            enable_thinking=enable_thinking,
-        )
-        inputs = processor(
-            text=[text],
-            return_tensors="pt",
-            padding=True,
-        ).to("cuda")
-
-    # Stream the response
-    streamer = TextIteratorStreamer(
-        processor.tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True,
-        timeout=120,
-    )
+        user_content.append({"type": "image", "image": image})
+    user_content.append({"type": "text", "text": message if message.strip() else "Describe this image."})
+    messages.append({"role": "user", "content": user_content})
 
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        use_cache=True,
+    text = processor.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=enable_thinking,
     )
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    # Collect images from messages
+    image_inputs = []
+    for msg in messages:
+        for item in msg.get("content", []):
+            if isinstance(item, dict) and item.get("type") == "image":
+                from PIL import Image
+                image_inputs.append(Image.open(item["image"]).convert("RGB"))
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs if image_inputs else None,
+        padding=True,
+        return_tensors="pt",
+    ).to("cuda")
 
-    full_response = ""
-    new_history = history + [(message, "")]
-    for token in streamer:
-        full_response += token
-        new_history[-1] = (message, full_response)
-        yield new_history, ""
+    do_sample = temperature > 0.001
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=do_sample,
+            temperature=float(temperature) if do_sample else 1.0,
+            top_p=float(top_p) if do_sample else 1.0,
+        )
 
-    thread.join()
-    yield new_history, ""
+    generated = output_ids[0][inputs["input_ids"].shape[1]:]
+    result = processor.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    return result
 
 
 CSS = """
@@ -134,86 +185,72 @@ CSS = """
 with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
     gr.Markdown(
         """
-        # 🧠 Grug-9B Reasoning VLM
-        A 9B-parameter vision-language reasoning model fine-tuned from Ornith-1.0-9B (Qwen3.5 architecture)
-        to "think small" — shorter internal reasoning while maintaining coding and agent capabilities.
+        # 🪨 Grug-9B Vision-Language Demo
+
+        **ProCreations/grug-9b** — a 9B-parameter reasoning VLM (fine-tuned from Ornith-1.0-9B / Qwen3.5)
+        that "thinks small" — producing concise reasoning instead of verbose chain-of-thought.
+
+        Upload an image and ask a question, or just type a prompt. The model will respond
+        with a short reasoning trace followed by its answer.
 
-        Upload an image (optional) and ask a question. The model reasons inside `` tags before answering.
+        [Model card](https://huggingface.co/ProCreations/grug-9b) · [Base model](https://huggingface.co/deepreinforce-ai/Ornith-1.0-9B)
         """
     )
 
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
+    with gr.Row():
+        with gr.Column(scale=3):
             image_input = gr.Image(
-                label="Image (optional)",
-                type="pil",
-                height=350,
+                label="Upload Image (optional)",
+                type="filepath",
+                height=320,
+            )
+            text_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Ask something about the image, or type a prompt…",
+                lines=3,
+            )
+            submit_btn = gr.Button("Submit", variant="primary")
+
+        with gr.Column(scale=4):
+            output_text = gr.Textbox(
+                label="Response",
+                lines=16,
+                max_lines=30,
+                show_copy_button=True,
             )
-            with gr.Column(scale=2):
-                chatbot = gr.Chatbot(
-                    label="Chat",
-                    height=450,
-                    show_label=True,
-                )
-                msg_input = gr.Textbox(
-                    label="Message",
-                    placeholder="Ask something about the image, or just chat...",
-                    show_label=False,
-                    lines=2,
-                )
-                with gr.Row():
-                    send_btn = gr.Button("Send", variant="primary", scale=2)
-                    clear_btn = gr.Button("Clear", scale=1)
-
-        with gr.Accordion("Advanced settings", open=False):
-            with gr.Row():
-                max_tokens = gr.Slider(
-                    label="Max new tokens",
-                    minimum=128,
-                    maximum=4096,
-                    value=1024,
-                    step=128,
-                )
-                temperature = gr.Slider(
-                    label="Temperature",
-                    minimum=0.1,
-                    maximum=2.0,
-                    value=0.7,
-                    step=0.1,
-                )
-                thinking_toggle = gr.Checkbox(
-                    label="Enable thinking (reasoning mode)",
-                    value=True,
-                )
-
-        gr.Examples(
-            examples=[
-                [None, "Write a Python function that checks if a number is prime.", 512, 0.7, True],
-                [None, "Explain the difference between TCP and UDP in networking.", 1024, 0.7, True],
-                ["examples/cat_tabby.jpg", "What breed is this cat? Describe what you see in detail.", 1024, 0.7, True],
-                ["examples/city_skyline_night.jpg", "What city might this be? Describe the architectural style.", 1024, 0.7, True],
-                ["examples/sushi_platter.jpg", "What kinds of sushi are on this platter?", 1024, 0.7, True],
-            ],
-            inputs=[image_input, msg_input, max_tokens, temperature, thinking_toggle],
-            outputs=[chatbot, msg_input],
-            fn=chat_with_image,
-            cache_examples=True,
-            cache_mode="lazy",
-        )
 
-    def clear_chat():
-        return [], ""
+    with gr.Accordion("Advanced settings", open=False):
+        with gr.Row():
+            max_tokens = gr.Slider(
+                label="Max new tokens", minimum=64, maximum=2048, value=512, step=64,
+            )
+            temperature = gr.Slider(
+                label="Temperature", minimum=0.0, maximum=2.0, value=0.7, step=0.05,
+            )
+            top_p = gr.Slider(
+                label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.05,
+            )
+            thinking = gr.Checkbox(label="Enable thinking (<think> tag)", value=True)
 
-    clear_btn.click(clear_chat, outputs=[chatbot, msg_input])
-    send_btn.click(
-        chat_with_image,
-        inputs=[image_input, msg_input, chatbot, max_tokens, temperature, thinking_toggle],
-        outputs=[chatbot, msg_input],
+    gr.Examples(
+        examples=[
+            ["examples/astronaut.jpg", "What is happening in this image?", 512, 0.7, 0.9, True],
+            ["examples/cat_tabby.jpg", "Describe this cat in detail.", 512, 0.7, 0.9, True],
+            ["examples/bird_bee_eater.jpg", "What species is this bird? What is it doing?", 512, 0.7, 0.9, True],
+        ],
+        inputs=[image_input, text_input, max_tokens, temperature, top_p, thinking],
+        outputs=output_text,
+        fn=predict,
+        cache_examples=True,
+        cache_mode="lazy",
     )
-    msg_input.submit(
-        chat_with_image,
-        inputs=[image_input, msg_input, chatbot, max_tokens, temperature, thinking_toggle],
-        outputs=[chatbot, msg_input],
+
+    submit_btn.click(
+        fn=predict,
+        inputs=[image_input, text_input, max_tokens, temperature, top_p, thinking],
+        outputs=output_text,
+        api_name="predict",
     )
 
-demo.launch(mcp_server=True)
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -1,8 +1,8 @@
-transformers==5.13.0
+transformers>=5.8
 accelerate
+torch
+torchvision
+Pillow
 sentencepiece
-qwen-vl-utils
 einops
-pillow
-numpy
-torchvision
+numpy