Switch assistant to gemma-4-12B (8-bit) + expose system prompt

- VLM_MODEL_ID -> OpenYourMind/gemma-4-12B-it-abliterated-uncensored
(model_type gemma4_unified). Load 8-bit via bitsandbytes by default
(~12GB, fits the zero-a10g alongside the diffusion pipelines);
VLM_LOAD_8BIT=0 reverts to bf16. Prior models kept commented.
- requirements: transformers -> git main (gemma4_unified support) + bitsandbytes.
- Expose an optional system_prompt through _vlm_chat_core (real system role
when the template supports it, folded into the user turn otherwise) and
surface it in the Prompt Assistant UI + the prompt_assistant /
prompt_to_video_assets / image_to_video_assets API endpoints.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

app.py

@@ -20,7 +20,12 @@ from diffusers import (
 )
 from compel import Compel, ReturnedEmbeddingsType
 from huggingface_hub import hf_hub_download
-from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
+from transformers import (
+    AutoProcessor,
+    AutoModelForImageTextToText,
+    BitsAndBytesConfig,
+    TextIteratorStreamer,
+)
 
 # Structured progress contract (also see wan2.2 / LTX2.3 Spaces and the generator
 # orchestrator). Every streaming endpoint yields a hidden JSON "progress" output
@@ -145,26 +150,36 @@ print("Pipelines loaded!")
 # describing a reference image. Loaded once, bf16 — fits alongside the two
 # diffusion pipelines on the shared GPU.
 # -----------------------------------------------------------------------------
-# Prior models, kept for easy revert:
+# Prior models, kept for easy revert (set VLM_MODEL_ID back to one of these):
 #   "rodrigomt/Qwen3.5-4B-Uncensored-Aggressive"   # no generation_config; needed eos pinning
-#   "prithivMLmods/gemma-4-E4B-it-Uncensored-MAX"  # under-rated explicit content (rating capped ~2)
-# Current: 9B uncensored Qwen3.5 VL. Vision-capable (image-text-to-text). NOTE its
-# generation_config eos is just <|endoftext|> and omits the chat terminator
-# <|im_end|>; _resolve_vlm_eos_ids() unions both so stopping still works. ~18.8 GB
-# bf16 — heaviest assistant so far; watch GPU memory alongside the diffusers.
-VLM_MODEL_ID = "ccharnkij/Qwen3.5-9B-Uncensored"
-
-print(f"Loading assistant: {VLM_MODEL_ID} ...")
+#   "prithivMLmods/gemma-4-E4B-it-Uncensored-MAX"  # model_type gemma4; under-rated explicit content (~2)
+#   "ccharnkij/Qwen3.5-9B-Uncensored"              # 9B Qwen3.5 VL, ~18.8 GB bf16 (the previous default)
+# Current: gemma-4-12B "unified" (text+image+audio+video), abliterated/uncensored.
+# model_type=gemma4_unified / Gemma4UnifiedForConditionalGeneration — NEWER than the
+# earlier gemma4, so it needs a very recent transformers (requirements pins git main).
+# It's ~24 GB in bf16, which does NOT fit on the zero-a10g (~24 GB) next to the two
+# diffusion pipelines, so we load it 8-bit (bitsandbytes, ~12 GB) by DEFAULT.
+# Flip VLM_LOAD_8BIT=0 to load full bf16 (only safe on a bigger GPU). Stop tokens are
+# handled model-agnostically by _resolve_vlm_eos_ids() (this model ships a proper
+# generation_config + chat_template.jinja).
+VLM_MODEL_ID = os.environ.get("VLM_MODEL_ID", "OpenYourMind/gemma-4-12B-it-abliterated-uncensored")
+VLM_LOAD_8BIT = os.environ.get("VLM_LOAD_8BIT", "1").lower() not in ("0", "false", "no", "")
+
+print(f"Loading assistant: {VLM_MODEL_ID} (8bit={VLM_LOAD_8BIT}) ...")
 vlm_processor = AutoProcessor.from_pretrained(
     VLM_MODEL_ID, token=HF_TOKEN, trust_remote_code=True
 )
-vlm_model = AutoModelForImageTextToText.from_pretrained(
-    VLM_MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    token=HF_TOKEN,
-    trust_remote_code=True,
-)
-vlm_model.to("cuda").eval()
+_vlm_kwargs = dict(token=HF_TOKEN, trust_remote_code=True, torch_dtype=torch.bfloat16)
+if VLM_LOAD_8BIT:
+    # bitsandbytes quantizes the weights on GPU as they load; device_map places them
+    # there directly, so we must NOT call .to("cuda") afterwards (it errors on a
+    # quantized model).
+    _vlm_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+    _vlm_kwargs["device_map"] = "auto"
+vlm_model = AutoModelForImageTextToText.from_pretrained(VLM_MODEL_ID, **_vlm_kwargs)
+if not VLM_LOAD_8BIT:
+    vlm_model.to("cuda")
+vlm_model.eval()
 
 
 def _resolve_vlm_eos_ids(tokenizer, model):
@@ -671,15 +686,21 @@ def _generate_image_inner(
 # =============================================================================
 # Prompt Assistant (Qwen3.5-4B) — single-turn chat, optional image
 # =============================================================================
-def _vlm_chat_core(message, image, reasoning, max_new_tokens):
+def _vlm_chat_core(message, image, reasoning, max_new_tokens, system_prompt=""):
     """Token-streaming VLM chat body (no GPU decorator).
 
     Yields ``("progress", produced, budget, partial)`` per token and a final
     ``("text", answer)`` tuple. Kept decorator-free so it can be reused inside a
     *single* ``@spaces.GPU`` window by other endpoints (e.g. the image→video
     check), avoiding a second ZeroGPU acquisition. See :func:`vlm_chat`.
+
+    ``system_prompt`` (optional) steers the assistant. Models whose chat template
+    supports a ``system`` role (e.g. Qwen) get a real system turn; templates that
+    reject one (e.g. Gemma) fall back to folding the system text into the user
+    turn, so the param works regardless of which model is loaded.
     """
     message = (message or "").strip()
+    system_prompt = (system_prompt or "").strip()
     if not message and image is None:
         yield ("text", "Please enter a question (and optionally attach an image).")
         return
@@ -688,20 +709,38 @@ def _vlm_chat_core(message, image, reasoning, max_new_tokens):
     budget = int(max_new_tokens)
     _gpu_start = time.time()
 
+    user_text = message or "Describe this image."
     content = []
     if image is not None:
         content.append({"type": "image", "image": image})
-    content.append({"type": "text", "text": message or "Describe this image."})
-    messages = [{"role": "user", "content": content}]
+    content.append({"type": "text", "text": user_text})
+
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
+    messages.append({"role": "user", "content": content})
+
+    def _template(msgs):
+        return vlm_processor.apply_chat_template(
+            msgs,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+            enable_thinking=enable_thinking,
+        )
 
-    inputs = vlm_processor.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        enable_thinking=enable_thinking,
-    ).to(vlm_model.device)
+    try:
+        inputs = _template(messages).to(vlm_model.device)
+    except Exception:
+        # Template has no system role (Gemma et al.) — fold it into the user turn.
+        if not system_prompt:
+            raise
+        folded = []
+        if image is not None:
+            folded.append({"type": "image", "image": image})
+        folded.append({"type": "text", "text": f"{system_prompt}\n\n{user_text}"})
+        inputs = _template([{"role": "user", "content": folded}]).to(vlm_model.device)
 
     tokenizer = getattr(vlm_processor, "tokenizer", vlm_processor)
     streamer = TextIteratorStreamer(
@@ -759,19 +798,20 @@ def _vlm_chat_core(message, image, reasoning, max_new_tokens):
 
 
 @spaces.GPU
-def vlm_chat(message, image, reasoning, max_new_tokens, progress=gr.Progress(track_tqdm=True)):
+def vlm_chat(message, image, reasoning, max_new_tokens, system_prompt="", progress=gr.Progress(track_tqdm=True)):
     """Answer a single user message, optionally grounded on an uploaded image.
 
     Thin ``@spaces.GPU`` wrapper around :func:`_vlm_chat_core` (which holds the
-    token-streaming logic). ``reasoning`` ("On"/"Off") drives Qwen's
+    token-streaming logic). ``reasoning`` ("On"/"Off") drives the model's
     ``enable_thinking`` switch: Off skips the <think> trace for a direct answer
     (best for prompt rewriting); On lets the model reason first (slower).
+    ``system_prompt`` (optional) steers the assistant.
     """
-    yield from _vlm_chat_core(message, image, reasoning, max_new_tokens)
+    yield from _vlm_chat_core(message, image, reasoning, max_new_tokens, system_prompt)
 
 
 @spaces.GPU
-def _image_video_check_gpu(image, message, reasoning, max_new_tokens):
+def _image_video_check_gpu(image, message, reasoning, max_new_tokens, system_prompt=""):
     """One GPU window: moderate ``image``, then stream a video prompt from it.
 
     Yields ``("moderation", dict)`` once the image-context check is done, then
@@ -785,7 +825,7 @@ def _image_video_check_gpu(image, message, reasoning, max_new_tokens):
         moderation = {"ok": False, "rating": None, "confidence": None,
                       "flags": None, "raw": "", "error": f"{type(exc).__name__}: {exc}"}
     yield ("moderation", moderation)
-    yield from _vlm_chat_core(message, image, reasoning, max_new_tokens)
+    yield from _vlm_chat_core(message, image, reasoning, max_new_tokens, system_prompt)
 
 
 def generate_and_upload(
@@ -869,16 +909,17 @@ def generate_and_upload(
 
 
 def assistant_chat(
-    message, image, reasoning, max_new_tokens,
+    message, image, reasoning, max_new_tokens, system_prompt="",
     progress=gr.Progress(track_tqdm=True),
 ):
     """Gradio-facing wrapper around ``vlm_chat``.
 
     Yields ``(answer, progress)`` so the UI streams the text live and a downstream
     consumer reading the progress index sees this node advance over SSE. The final
-    ``complete`` frame carries the clean answer at index 0.
+    ``complete`` frame carries the clean answer at index 0. ``system_prompt``
+    (optional) steers the assistant.
     """
-    for ev in vlm_chat(message, image, reasoning, max_new_tokens, progress=progress):
+    for ev in vlm_chat(message, image, reasoning, max_new_tokens, system_prompt, progress=progress):
         if ev[0] == "progress":
             _, produced, budget, partial = ev
             frac = min(0.99, produced / max(budget, 1))
@@ -934,6 +975,7 @@ def prompt_to_video_assets(
     randomize_seed: bool = True,
     reasoning: str = "Off",
     max_new_tokens: int = 512,
+    system_prompt: str = "",
     request: gr.Request = None,
 ) -> dict:
     """Text -> first-frame image (uploaded to R2) + video generation prompt.
@@ -977,7 +1019,7 @@ def prompt_to_video_assets(
     frame_prompt = ""
     for ev in vlm_chat(
         _compose_instruction(image_instruction, original_input),
-        None, reasoning, max_new_tokens,
+        None, reasoning, max_new_tokens, system_prompt,
     ):
         if ev[0] == "progress":
             _, produced, budget, partial = ev
@@ -1044,7 +1086,7 @@ def prompt_to_video_assets(
     video_prompt = ""
     for ev in vlm_chat(
         _compose_instruction(video_instruction, original_input),
-        image, reasoning, max_new_tokens,
+        image, reasoning, max_new_tokens, system_prompt,
     ):
         if ev[0] == "progress":
             _, produced, budget, partial = ev
@@ -1088,6 +1130,7 @@ def image_to_video_assets(
     video_instruction,
     reasoning,
     max_new_tokens,
+    system_prompt="",
     progress=gr.Progress(track_tqdm=True),
 ):
     """User-uploaded image -> motion prompt + image-context (moderation) check.
@@ -1110,7 +1153,7 @@ def image_to_video_assets(
     video_prompt = ""
     yield video_prompt, status, _progress("check", 0.02, label="Screening image")
 
-    for ev in _image_video_check_gpu(image, instruction, reasoning, budget):
+    for ev in _image_video_check_gpu(image, instruction, reasoning, budget, system_prompt):
         if ev[0] == "moderation":
             status = {"moderation": ev[1]}
             yield video_prompt, status, _progress("check", 0.2, label="Image checked")
@@ -1333,9 +1376,10 @@ with gr.Blocks(fill_height=True) as demo:
 
         with gr.Tab("💬 Prompt Assistant"):
             gr.Markdown(
-                "Ask **Qwen3.5-9B** a single question — with or without an image. "
+                "Ask the assistant model a single question — with or without an image. "
                 "Turn a rough idea into a rich prompt, or upload a reference image and "
-                "ask the model to describe it as a prompt."
+                "ask the model to describe it as a prompt. Use the **System Prompt** to "
+                "steer its tone/role."
             )
             with gr.Row(equal_height=False):
                 with gr.Column(scale=1, min_width=320):
@@ -1344,6 +1388,12 @@ with gr.Blocks(fill_height=True) as demo:
                         type="pil",
                         height=320,
                     )
+                    vlm_system = gr.Textbox(
+                        label="🧭 System Prompt (optional)",
+                        placeholder="e.g. 'You are an expert Danbooru tag writer. Output only comma-separated tags.'",
+                        lines=2,
+                        max_lines=8,
+                    )
                     vlm_prompt = gr.Textbox(
                         label="💬 Your Question",
                         placeholder="e.g. 'Improve this prompt: a cat in a hat' — or attach an image and ask 'describe this as a prompt'",
@@ -1413,8 +1463,8 @@ with gr.Blocks(fill_height=True) as demo:
         outputs=[output_image, used_seed, r2_status, gen_progress],
     )
 
-    # Prompt Assistant (Qwen3.5-4B) — single-turn, optional image
-    vlm_inputs = [vlm_prompt, vlm_image, vlm_reasoning, vlm_max_tokens]
+    # Prompt Assistant — single-turn, optional image, optional system prompt
+    vlm_inputs = [vlm_prompt, vlm_image, vlm_reasoning, vlm_max_tokens, vlm_system]
     vlm_btn.click(
         fn=assistant_chat, inputs=vlm_inputs, outputs=[vlm_output, vlm_progress],
         api_name="prompt_assistant",
@@ -1432,13 +1482,14 @@ with gr.Blocks(fill_height=True) as demo:
         i2v_instruction = gr.Textbox(label="i2v instruction", value="")
         i2v_reasoning = gr.Radio(choices=["Off", "On"], value="Off", label="i2v reasoning")
         i2v_max_tokens = gr.Slider(minimum=64, maximum=1024, value=256, step=64, label="i2v max tokens")
+        i2v_system = gr.Textbox(label="i2v system prompt", value="")
         i2v_btn = gr.Button("i2v", visible=False)
         i2v_prompt_out = gr.Textbox(label="i2v prompt")
         i2v_status_out = gr.JSON(label="i2v status")
         i2v_progress_out = gr.JSON(label="i2v progress")
     i2v_btn.click(
         fn=image_to_video_assets,
-        inputs=[i2v_image, i2v_instruction, i2v_reasoning, i2v_max_tokens],
+        inputs=[i2v_image, i2v_instruction, i2v_reasoning, i2v_max_tokens, i2v_system],
         outputs=[i2v_prompt_out, i2v_status_out, i2v_progress_out],
         api_name="image_to_video_assets",
     )

requirements.txt

@@ -1,8 +1,10 @@
 gradio[mcp]
 git+https://github.com/huggingface/diffusers
-transformers
+# gemma-4-12B is model_type gemma4_unified — needs a very recent transformers.
+git+https://github.com/huggingface/transformers
 torchvision
 accelerate
+bitsandbytes
 compel
 sentencepiece
 numpy