feat: add Qwen3.5-4B Prompt Assistant tab (text + image, single-turn)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -16,6 +16,7 @@ from diffusers import (
 )
 from compel import Compel, ReturnedEmbeddingsType
 from huggingface_hub import hf_hub_download
+from transformers import AutoProcessor, AutoModelForImageTextToText
 
 import r2_uploader
 
@@ -107,6 +108,28 @@ noobxl_pipe.to("cuda")
 
 print("Pipelines loaded!")
 
+# -----------------------------------------------------------------------------
+# Qwen3.5-4B multimodal assistant. Powers the "Prompt Assistant" tab: one-turn
+# question answering with an optional image, for turning rough ideas into rich
+# prompts or describing a reference image. Loaded once, bf16 (~9-10 GB) — fits
+# alongside the two diffusion pipelines on the shared GPU.
+# -----------------------------------------------------------------------------
+VLM_MODEL_ID = "rodrigomt/Qwen3.5-4B-Uncensored-Aggressive"
+
+print("Loading Qwen3.5-4B assistant...")
+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()
+
+print("Assistant loaded!")
+
 
 # =============================================================================
 # Compel-based prompt weighting helpers (ported from the NoobXL11 reference
@@ -409,6 +432,50 @@ def _generate_image_inner(
     return image, seed
 
 
+# =============================================================================
+# Prompt Assistant (Qwen3.5-4B) — single-turn chat, optional image
+# =============================================================================
+@spaces.GPU
+def vlm_chat(message, image, max_new_tokens, progress=gr.Progress(track_tqdm=True)):
+    """Answer a single user message, optionally grounded on an uploaded image."""
+    message = (message or "").strip()
+    if not message and image is None:
+        return "Please enter a question (and optionally attach an image)."
+
+    _gpu_start = time.time()
+    try:
+        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}]
+
+        inputs = vlm_processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(vlm_model.device)
+
+        with torch.inference_mode():
+            generated = vlm_model.generate(
+                **inputs,
+                max_new_tokens=int(max_new_tokens),
+                do_sample=False,
+            )
+        # Drop the prompt tokens so only the freshly generated answer is decoded.
+        trimmed = generated[0][inputs["input_ids"].shape[1]:]
+        text = vlm_processor.decode(trimmed, skip_special_tokens=True)
+        return text.strip()
+    finally:
+        print(
+            f"[ImageStudio] Assistant GPU time: {time.time() - _gpu_start:.2f}s "
+            f"(has_image={image is not None}, max_new_tokens={int(max_new_tokens)})",
+            flush=True,
+        )
+
+
 def generate_and_upload(
     model_name,
     prompt,
@@ -560,90 +627,127 @@ with gr.Blocks(fill_height=True) as demo:
         elem_classes="header-text"
     )
 
-    with gr.Row(equal_height=False):
-        # Left column - Input controls
-        with gr.Column(scale=1, min_width=320):
-            model_name = gr.Dropdown(
-                choices=MODEL_CHOICES,
-                value=MODEL_ZIMAGE,
-                label="🧠 Model",
-                info="Moody Pro Mix: fast & general (Z-Image) • One Obsession: anime / illustration (SDXL)",
-            )
-
-            prompt = gr.Textbox(
-                label="✨ Your Prompt",
-                placeholder="Describe the image you want to create...",
-                lines=5,
-                max_lines=10,
-                autofocus=True,
-            )
+    with gr.Tabs():
+        with gr.Tab("🎨 Generate"):
+            with gr.Row(equal_height=False):
+                # Left column - Input controls
+                with gr.Column(scale=1, min_width=320):
+                    model_name = gr.Dropdown(
+                        choices=MODEL_CHOICES,
+                        value=MODEL_ZIMAGE,
+                        label="🧠 Model",
+                        info="Moody Pro Mix: fast & general (Z-Image) • One Obsession: anime / illustration (SDXL)",
+                    )
 
-            with gr.Accordion("⚙️ Advanced Settings", open=False):
-                use_negative_prompt = gr.Checkbox(
-                    label="Use Negative Prompt",
-                    value=True,
-                    visible=False,  # NoobXL only
-                )
-                negative_prompt = gr.Textbox(
-                    label="🚫 Negative Prompt",
-                    value=NOOBXL_NEGATIVE,
-                    lines=3,
-                    max_lines=6,
-                    visible=False,  # NoobXL only
-                )
-
-                with gr.Row():
-                    height = gr.Slider(
-                        minimum=512, maximum=2048, value=1024, step=64,
-                        label="Height", info="Image height in pixels",
+                    prompt = gr.Textbox(
+                        label="✨ Your Prompt",
+                        placeholder="Describe the image you want to create...",
+                        lines=5,
+                        max_lines=10,
+                        autofocus=True,
                     )
-                    width = gr.Slider(
-                        minimum=512, maximum=2048, value=1024, step=64,
-                        label="Width", info="Image width in pixels",
+
+                    with gr.Accordion("⚙️ Advanced Settings", open=False):
+                        use_negative_prompt = gr.Checkbox(
+                            label="Use Negative Prompt",
+                            value=True,
+                            visible=False,  # NoobXL only
+                        )
+                        negative_prompt = gr.Textbox(
+                            label="🚫 Negative Prompt",
+                            value=NOOBXL_NEGATIVE,
+                            lines=3,
+                            max_lines=6,
+                            visible=False,  # NoobXL only
+                        )
+
+                        with gr.Row():
+                            height = gr.Slider(
+                                minimum=512, maximum=2048, value=1024, step=64,
+                                label="Height", info="Image height in pixels",
+                            )
+                            width = gr.Slider(
+                                minimum=512, maximum=2048, value=1024, step=64,
+                                label="Width", info="Image width in pixels",
+                            )
+
+                        num_inference_steps = gr.Slider(
+                            minimum=1, maximum=50, value=9, step=1,
+                            label="Inference Steps",
+                            info="Moody Pro Mix: ~9 • One Obsession: ~28",
+                        )
+
+                        guidance_scale = gr.Slider(
+                            minimum=0.0, maximum=10.0, value=0.0, step=0.1,
+                            label="Guidance Scale",
+                            info="Moody Pro Mix (Z-Image) is guidance-free (forced to 0)",
+                            interactive=False,
+                        )
+
+                        with gr.Row():
+                            randomize_seed = gr.Checkbox(label="🎲 Random Seed", value=True)
+                            seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED)
+
+                    generate_btn = gr.Button(
+                        "🚀 Generate Image", variant="primary", size="lg", scale=1
                     )
 
-                num_inference_steps = gr.Slider(
-                    minimum=1, maximum=50, value=9, step=1,
-                    label="Inference Steps",
-                    info="Moody Pro Mix: ~9 • One Obsession: ~28",
-                )
-
-                guidance_scale = gr.Slider(
-                    minimum=0.0, maximum=10.0, value=0.0, step=0.1,
-                    label="Guidance Scale",
-                    info="Moody Pro Mix (Z-Image) is guidance-free (forced to 0)",
-                    interactive=False,
-                )
-
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(label="🎲 Random Seed", value=True)
-                    seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED)
-
-            generate_btn = gr.Button(
-                "🚀 Generate Image", variant="primary", size="lg", scale=1
-            )
+                    gr.Examples(
+                        examples=examples,
+                        inputs=[prompt],
+                        label="💡 Try these prompts",
+                        examples_per_page=5,
+                    )
 
-            gr.Examples(
-                examples=examples,
-                inputs=[prompt],
-                label="💡 Try these prompts",
-                examples_per_page=5,
-            )
+                # Right column - Output
+                with gr.Column(scale=1, min_width=320):
+                    output_image = gr.Image(
+                        label="Generated Image",
+                        type="pil",
+                        format="png",
+                        show_label=False,
+                        height=600,
+                        buttons=["download", "share"],
+                    )
+                    used_seed = gr.Number(
+                        label="🎲 Seed Used", interactive=False, container=True,
+                    )
+                    r2_status = gr.JSON(label="☁️ R2 Upload")
 
-        # Right column - Output
-        with gr.Column(scale=1, min_width=320):
-            output_image = gr.Image(
-                label="Generated Image",
-                type="pil",
-                format="png",
-                show_label=False,
-                height=600,
-                buttons=["download", "share"],
-            )
-            used_seed = gr.Number(
-                label="🎲 Seed Used", interactive=False, container=True,
+        with gr.Tab("💬 Prompt Assistant"):
+            gr.Markdown(
+                "Ask **Qwen3.5-4B** 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."
             )
-            r2_status = gr.JSON(label="☁️ R2 Upload")
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=1, min_width=320):
+                    vlm_image = gr.Image(
+                        label="🖼️ Image (optional)",
+                        type="pil",
+                        height=320,
+                    )
+                    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'",
+                        lines=4,
+                        max_lines=12,
+                    )
+                    with gr.Accordion("⚙️ Settings", open=False):
+                        vlm_max_tokens = gr.Slider(
+                            minimum=64, maximum=2048, value=512, step=64,
+                            label="Max New Tokens",
+                            info="Upper bound on the answer length",
+                        )
+                    vlm_btn = gr.Button(
+                        "✨ Ask", variant="primary", size="lg", scale=1
+                    )
+                with gr.Column(scale=1, min_width=320):
+                    vlm_output = gr.Textbox(
+                        label="🤖 Answer",
+                        lines=20,
+                        show_copy_button=True,
+                    )
 
     gr.Markdown(
         """
@@ -682,6 +786,14 @@ with gr.Blocks(fill_height=True) as demo:
         outputs=[output_image, used_seed, r2_status],
     )
 
+    # Prompt Assistant (Qwen3.5-4B) — single-turn, optional image
+    vlm_inputs = [vlm_prompt, vlm_image, vlm_max_tokens]
+    vlm_btn.click(
+        fn=vlm_chat, inputs=vlm_inputs, outputs=[vlm_output],
+        api_name="prompt_assistant",
+    )
+    vlm_prompt.submit(fn=vlm_chat, inputs=vlm_inputs, outputs=[vlm_output])
+
 if __name__ == "__main__":
     demo.launch(
         theme=custom_theme,