Add a prompt upsample ui for upsampling prompts to ideogram format prompts.

extensions_built_in/captioner/prompts/ideogram4_upsample_prompt.py

@@ -0,0 +1,83 @@
+ideogram4_upsample_prompt = """
+[META]
+frozen: false
+description: Faithful upsampler — lays a user prompt into the structured JSON caption without inventing or embellishing. Preserves triggers/names/styles exactly. Thinking off.
+thinking_mode: disabled
+
+[SYSTEM]
+You convert a user prompt into a structured JSON caption an image renderer can consume. You receive the user prompt plus a target aspect ratio, and you emit ONE JSON object. Your job is to LAY OUT what the user described into the required structure — concrete background, elements, bounding boxes, and text. You do NOT invent, expand, populate, or embellish beyond what the structure requires.
+
+## FIDELITY — read first, applies above everything else
+
+- **Preserve triggers/tokens EXACTLY.** Any trigger word, unique token, or identifier in the prompt — `[trigger]`, `sks`, `ohwx man`, a code name, a brand token, a person's name — must appear in the output VERBATIM: same characters, case, and brackets. Never paraphrase, translate, pluralize, split, correct, or drop it. Put it in the `desc` (and `high_level_description`) of the element it refers to.
+- **Named person → no invented appearance.** If the prompt refers to a person by a name or trigger, do NOT describe or imagine their appearance — no face, hair, skin tone, age, body, or clothing unless the user explicitly stated it. Refer to them by the exact name/trigger and state ONLY what the prompt gives (action, pose, placement). Their identity is carried by the name alone.
+- **Named style → no invented style detail.** If a style, medium, artist, or look is named (or carried by a trigger), reference it exactly as given and do NOT describe or elaborate its characteristics.
+{{mode_directive}}
+
+## OUTPUT CONTRACT — exactly three top-level keys, in this order:
+
+```json
+{"aspect_ratio":"W:H","high_level_description":"...","compositional_deconstruction":{"background":"...","elements":[ ... ]}}
+```
+
+- Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys.
+- Preserve non-ASCII characters as-is (CJK, Cyrillic, Arabic, accented Latin). Never escape them as unicode code-point sequences or transliterate.
+- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`). The `text` field is the exception — it holds verbatim characters.
+
+### `aspect_ratio` (first field)
+
+The target ratio is given. Echo it VERBATIM. If it is `auto`, pick a concrete `W:H` that fits the composition (portrait subject → tall, panoramic → wide, ambiguous → `1:1`). Never emit `auto`.
+
+### `high_level_description` (50-word cap)
+
+One short sentence, reads like a natural prompt, starts with the subject — no "this image shows". Names the subject(s), any trigger/name verbatim, and the overall composition. Don't enumerate fine detail.
+
+## ELEMENTS
+
+Each element is one of:
+```
+{"type":"obj","bbox":[y1,x1,y2,x2],"desc":"..."}
+{"type":"text","bbox":[y1,x1,y2,x2],"text":"LINE ONE\nLINE TWO","desc":"..."}
+```
+`bbox` is optional per element (see BBOX).
+
+- **One coherent subject = ONE element.** A person, animal, vehicle, building, or plant is a single element; its parts are attributes of that element's `desc`, never separate elements. Multiple distinct subjects = multiple elements (one each).
+- **`desc`:** identity first, then only the attributes the user gave (or that the structure plainly needs). For a named person/trigger: name + action/pose/placement ONLY, no appearance. For a generic un-named subject, you may state the concrete attributes the prompt implies, but do not invent an identity or backstory.
+
+## BACKGROUND — the scene shell only
+
+`background` describes the shell: walls/finishes, floor/ground, sky, ambient light, and distant out-of-focus context.
+
+- The floor/ground/turf/pavement, sky, horizon, and distant crowds live in `background` ONLY — never as obj elements. (A floor emitted as an obj clips standing subjects' legs.)
+- **No double-counting:** anything named in `background` must NOT also be an obj element.
+- Don't smuggle furniture or people into `background` as a "receding arrangement" — those are foreground elements.
+- If the prompt asks for a transparent/cutout background, set `background` to exactly: `transparent background` (and include `on a transparent background` in the HLD).
+
+## BBOX
+
+Coordinates are normalized to 0–1000 in BOTH axes, top-left origin. Format `[y1, x1, y2, x2]` with `y1 < y2`, `x1 < x2`.
+
+A box is square only on a square frame; on a wide or tall frame the same numbers stretch. For round or square on-screen subjects, scale the spans so `(x2-x1)/(y2-y1) ≈ W/H`. Include bboxes where position matters; omit them for dense/uncountable fills (crowds, starfields).
+
+## TEXT
+
+- Every quoted string in the prompt becomes its own `text` element, with `text` = the verbatim characters (preserve case, punctuation, diacritics, and any trigger). Use `\n` for line breaks within one text block; separate blocks get separate elements.
+- Include clearly in-scene text (a sign, a label) only when the user asked for it — do not invent signage or brand copy.
+- Prose fields (`desc`, `background`, `high_level_description`) are always in ENGLISH; only the `text` field follows the prompt's language.
+
+## SPECIFICITY
+
+- For details the user GAVE, commit to one concrete value — no hedging (`things like`, `such as`, `various`), no alternatives (`oak or walnut`).
+- For details the user did NOT give, add a single concrete value only when the structure requires it (e.g. a plain background shell); otherwise leave it out.
+- Never hedge, never invent appearance for a named person, and never invent characteristics for a named style.
+
+## ADDITIONAL INSTRUCTIONS
+
+Honor the following extra instructions from the user. They must NEVER override the OUTPUT CONTRACT, the FIDELITY rules, or the structure above.
+
+{{user_instructions}}
+
+[USER]
+TARGET IMAGE ASPECT RATIO: {{aspect_ratio}} (width:height).
+User prompt: {{original_prompt}}
+"""

ui/src/app/jobs/new/SimpleJob.tsx

@@ -22,7 +22,8 @@ import {
   SliderInput,
 } from '@/components/formInputs';
 import Card from '@/components/Card';
-import { X, Copy } from 'lucide-react';
+import { X, Copy, Wand2 } from 'lucide-react';
+import { openUpsamplePromptsModal, toAspectRatio } from '@/components/UpsamplePromptsModal';
 import AddSingleImageModal, { openAddImageModal } from '@/components/AddSingleImageModal';
 import SampleControlImage from '@/components/SampleControlImage';
 import { FlipHorizontal2, FlipVertical2 } from 'lucide-react';
@@ -1263,9 +1264,35 @@ export default function SimpleJob({
                 </FormGroup>
               </div>
             </div>
-            <FormGroup label={`Sample Prompts (${jobConfig.config.process[0].sample.samples.length})`} className="pt-2">
-              <div></div>
-            </FormGroup>
+            <div className="pt-2 mb-2 flex items-center justify-between">
+              <label className="block text-xs text-gray-300">
+                Sample Prompts ({jobConfig.config.process[0].sample.samples.length})
+              </label>
+              {modelArch?.additionalSections?.includes('upsample_prompts') && (
+                <button
+                  type="button"
+                  disabled={jobConfig.config.process[0].sample.samples.length === 0}
+                  onClick={() => {
+                    const sampleCfg = jobConfig.config.process[0].sample;
+                    const items = sampleCfg.samples
+                      .map((s, i) => ({
+                        index: i,
+                        prompt: s.prompt || '',
+                        aspectRatio: toAspectRatio(s.width || sampleCfg.width, s.height || sampleCfg.height),
+                      }))
+                      .filter(it => it.prompt.trim() !== '');
+                    if (items.length === 0) return;
+                    openUpsamplePromptsModal(items, (index, newPrompt) => {
+                      setJobConfig(newPrompt, `config.process[0].sample.samples[${index}].prompt`);
+                    });
+                  }}
+                  className="px-3 py-1.5 text-sm bg-purple-600 hover:bg-purple-700 disabled:opacity-40 disabled:cursor-not-allowed text-white rounded-md inline-flex items-center gap-2"
+                >
+                  <Wand2 className="w-4 h-4" />
+                  Upsample Prompts
+                </button>
+              )}
+            </div>
             {jobConfig.config.process[0].sample.samples.map((sample, i) => (
               <div key={i} className="rounded-lg pl-4 pr-1 mb-4 bg-gray-950">
                 <div className="flex items-center space-x-2">

ui/src/app/jobs/new/options.ts

@@ -31,7 +31,8 @@ type AdditionalSections =
   | 'model.layer_offloading'
   | 'model.low_vram'
   | 'model.qie.match_target_res'
-  | 'model.assistant_lora_path';
+  | 'model.assistant_lora_path'
+  | 'upsample_prompts';
 
 type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
 
@@ -999,6 +1000,7 @@ export const modelArchs: ModelArch[] = [
     additionalSections: [
       'model.low_vram',
       'model.layer_offloading',
+      'upsample_prompts',
     ],
     hasMultiLinePrompts: true,
   },

ui/src/app/layout.tsx

@@ -10,6 +10,7 @@ import DocModal from '@/components/DocModal';
 import os from 'os';
 import { CaptionDatasetModal } from '@/components/CaptionDatasetModal';
 import MergeLoRAsModal from '@/components/MergeLoRAsModal';
+import UpsamplePromptsModal from '@/components/UpsamplePromptsModal';
 
 export const dynamic = 'force-dynamic';
 
@@ -63,6 +64,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
         <DocModal />
         <CaptionDatasetModal />
         <MergeLoRAsModal />
+        <UpsamplePromptsModal />
       </body>
     </html>
   );

ui/src/components/UpsamplePromptsModal.tsx

@@ -0,0 +1,362 @@
+'use client';
+import React, { useEffect, useRef, useState } from 'react';
+import { createGlobalState } from 'react-global-hooks';
+import { Check, X, Loader2, AlertTriangle } from 'lucide-react';
+import classNames from 'classnames';
+import { Modal } from './Modal';
+import { callScriptStream } from '@/utils/callScript';
+
+export interface UpsamplePromptItem {
+  index: number; // index into the samples array (write-back target)
+  prompt: string; // current prompt text
+  aspectRatio: string; // 'W:H' or 'auto'
+}
+
+export interface UpsamplePromptsModalState {
+  prompts: UpsamplePromptItem[];
+  onApply: (index: number, newPrompt: string) => void;
+  onClose?: () => void;
+}
+
+export const upsamplePromptsModalState = createGlobalState<UpsamplePromptsModalState | null>(null);
+
+export const openUpsamplePromptsModal = (
+  prompts: UpsamplePromptItem[],
+  onApply: (index: number, newPrompt: string) => void,
+  onClose?: () => void,
+) => {
+  upsamplePromptsModalState.set({ prompts, onApply, onClose });
+};
+
+// Reduce/snap pixel dimensions to a clean 'W:H' (denominator <= 16), mirroring the
+// captioner's compute_aspect_ratio. Returns 'auto' when dimensions are unknown.
+export function toAspectRatio(width?: number, height?: number): string {
+  if (!width || !height || width <= 0 || height <= 0) return 'auto';
+  const gcd = (a: number, b: number): number => (b ? gcd(b, a % b) : a);
+  const g = gcd(width, height);
+  const rw = width / g;
+  const rh = height / g;
+  const MAXD = 16;
+  if (rw <= MAXD && rh <= MAXD) return `${rw}:${rh}`;
+  const target = width / height;
+  let best: { err: number; p: number; q: number } | null = null;
+  for (let q = 1; q <= MAXD; q++) {
+    const p = Math.max(1, Math.round(target * q));
+    const err = Math.abs(p / q - target);
+    if (!best || err < best.err) best = { err, p, q };
+  }
+  return `${best!.p}:${best!.q}`;
+}
+
+type RowStatus = 'idle' | 'queued' | 'running' | 'done' | 'failed';
+
+const UpsamplePromptsModal: React.FC = () => {
+  const [modalInfo, setModalInfo] = upsamplePromptsModalState.use();
+  const isOpen = modalInfo !== null;
+
+  const [selected, setSelected] = useState<Record<number, boolean>>({});
+  const [creative, setCreative] = useState(false);
+  const [instructions, setInstructions] = useState('');
+  const [status, setStatus] = useState<Record<number, RowStatus>>({});
+  const [upsampled, setUpsampled] = useState<Record<number, string>>({});
+  const [isRunning, setIsRunning] = useState(false);
+  const [isDone, setIsDone] = useState(false);
+  const [hasError, setHasError] = useState(false);
+  const [logOutput, setLogOutput] = useState('');
+  const logRef = useRef<HTMLDivElement | null>(null);
+
+  // Reset state when the modal opens/closes. Nothing is selected by default.
+  useEffect(() => {
+    setSelected({});
+    setCreative(false);
+    setInstructions('');
+    setStatus({});
+    setUpsampled({});
+    setIsRunning(false);
+    setIsDone(false);
+    setHasError(false);
+    setLogOutput('');
+  }, [modalInfo]);
+
+  useEffect(() => {
+    if (logRef.current) logRef.current.scrollTop = logRef.current.scrollHeight;
+  }, [logOutput]);
+
+  const onClose = () => {
+    if (isRunning) return;
+    setModalInfo(null);
+    modalInfo?.onClose?.();
+  };
+
+  const prompts = modalInfo?.prompts ?? [];
+  const selectedCount = prompts.filter(p => selected[p.index]).length;
+
+  const toggle = (index: number) => {
+    if (isRunning) return;
+    setSelected(prev => ({ ...prev, [index]: !prev[index] }));
+  };
+
+  const setAll = (value: boolean) => {
+    if (isRunning) return;
+    const next: Record<number, boolean> = {};
+    prompts.forEach(p => (next[p.index] = value));
+    setSelected(next);
+  };
+
+  const onRun = async () => {
+    if (isRunning || !modalInfo) return;
+    const submitted = modalInfo.prompts.filter(p => selected[p.index]);
+    if (submitted.length === 0) return;
+
+    setIsRunning(true);
+    setIsDone(false);
+    setHasError(false);
+    setLogOutput('');
+    setUpsampled({});
+
+    const submittedIndices = submitted.map(p => p.index);
+    setStatus(() => {
+      const s: Record<number, RowStatus> = {};
+      submittedIndices.forEach((idx, k) => (s[idx] = k === 0 ? 'running' : 'queued'));
+      return s;
+    });
+
+    const payload = submitted.map(p => ({ prompt: p.prompt, aspect_ratio: p.aspectRatio }));
+
+    // The script streams one compact JSON line per completed prompt on stdout.
+    // stdout chunks don't align to line boundaries, so buffer and split ourselves.
+    let buffer = '';
+    const handleStdout = (chunk: string) => {
+      buffer += chunk;
+      let nl: number;
+      while ((nl = buffer.indexOf('\n')) >= 0) {
+        const line = buffer.slice(0, nl).trim();
+        buffer = buffer.slice(nl + 1);
+        if (!line) continue;
+        try {
+          const evt = JSON.parse(line);
+          if (typeof evt.index !== 'number') continue;
+          const orig = submittedIndices[evt.index];
+          if (orig === undefined) continue;
+          if (evt.caption) {
+            const text = JSON.stringify(evt.caption, null, 2);
+            setUpsampled(prev => ({ ...prev, [orig]: text }));
+            setStatus(prev => ({ ...prev, [orig]: 'done' }));
+            modalInfo.onApply(orig, text); // write back to the job config live
+          } else {
+            setStatus(prev => ({ ...prev, [orig]: 'failed' }));
+          }
+          const nextIdx = submittedIndices[evt.index + 1];
+          if (nextIdx !== undefined) {
+            setStatus(prev => (prev[nextIdx] === 'queued' ? { ...prev, [nextIdx]: 'running' } : prev));
+          }
+        } catch {
+          // ignore non-JSON / partial lines
+        }
+      }
+    };
+
+    const append = (chunk: string) => setLogOutput(prev => prev + chunk);
+
+    try {
+      const finalEvent = await callScriptStream('upsample_ideogram4_caption.py', {
+        args: {
+          prompts: JSON.stringify(payload),
+          stream: true,
+          temperature: 0.7,
+          quantize: true, // float8 (script default qtype) to reduce VRAM
+          ...(creative ? { creative: true } : {}),
+          ...(instructions.trim() ? { instructions: instructions.trim() } : {}),
+        },
+        onStdout: handleStdout,
+        onStderr: append,
+      });
+
+      const ok = finalEvent?.type === 'exit' && finalEvent.ok === true;
+      if (!ok) {
+        setHasError(true);
+        if (finalEvent?.type === 'error' && finalEvent.message) {
+          append(`\n${finalEvent.message}\n`);
+        } else if (finalEvent?.type === 'exit' && finalEvent.timedOut) {
+          append('\nScript timed out.\n');
+        } else if (finalEvent?.type === 'exit') {
+          append(`\nScript exited with code ${finalEvent.exitCode}.\n`);
+        }
+      }
+    } catch (err: any) {
+      setHasError(true);
+      append(`\n${err?.message || 'Unknown error'}\n`);
+    } finally {
+      // Anything still queued/running never reported a result -> failed.
+      setStatus(prev => {
+        const cp = { ...prev };
+        submittedIndices.forEach(i => {
+          if (cp[i] === 'running' || cp[i] === 'queued') cp[i] = 'failed';
+        });
+        return cp;
+      });
+      setIsRunning(false);
+      setIsDone(true);
+    }
+  };
+
+  const StatusIcon = ({ s }: { s: RowStatus | undefined }) => {
+    if (s === 'done') return <Check className="w-4 h-4 text-emerald-400" />;
+    if (s === 'failed') return <X className="w-4 h-4 text-rose-400" />;
+    if (s === 'running') return <Loader2 className="w-4 h-4 text-amber-400 animate-spin" />;
+    if (s === 'queued') return <Loader2 className="w-4 h-4 text-gray-500" />;
+    return <span className="inline-block w-4 h-4" />;
+  };
+
+  return (
+    <Modal
+      isOpen={isOpen}
+      onClose={onClose}
+      title="Upsample Prompts"
+      size="lg"
+      showCloseButton={!isRunning}
+      closeOnOverlayClick={!isRunning}
+    >
+      <div>
+        <div className="mb-3 flex items-start gap-2 rounded-md border border-amber-500/40 bg-amber-500/10 px-3 py-2 text-xs text-amber-300">
+          <AlertTriangle className="w-4 h-4 flex-shrink-0 mt-0.5" />
+          <span>
+            This loads a model and runs on the GPU. You need at least 13GB of free VRAM to run this, so make sure the
+            GPU is idle (no training or other jobs running) before starting, or it may run out of memory.
+          </span>
+        </div>
+
+        <div className="mb-2 flex items-center justify-between text-sm">
+          <div className="text-gray-400">
+            {isRunning && <span className="text-amber-400">Upsampling... please do not close this window.</span>}
+            {isDone && hasError && <span className="text-rose-400">Finished with errors. See log below.</span>}
+            {isDone && !hasError && <span className="text-emerald-400">Upsampling complete.</span>}
+            {!isRunning && !isDone && (
+              <span>
+                Select prompts to upsample into structured Ideogram captions. {selectedCount}/{prompts.length} selected.
+              </span>
+            )}
+          </div>
+          {!isRunning && !isDone && (
+            <div className="flex gap-2 flex-shrink-0">
+              <button type="button" onClick={() => setAll(true)} className="text-xs text-gray-300 hover:text-gray-100">
+                Select all
+              </button>
+              <span className="text-gray-600">|</span>
+              <button type="button" onClick={() => setAll(false)} className="text-xs text-gray-300 hover:text-gray-100">
+                None
+              </button>
+            </div>
+          )}
+        </div>
+
+        <div className="mb-2 flex items-center gap-2">
+          <button
+            type="button"
+            disabled={isRunning || isDone}
+            onClick={() => setCreative(c => !c)}
+            className={classNames(
+              'px-2.5 py-1 text-xs rounded-md border transition-colors disabled:opacity-40 disabled:cursor-not-allowed',
+              {
+                'bg-purple-600 border-purple-500 text-white': creative,
+                'border-gray-700 text-gray-300 hover:border-gray-500': !creative,
+              },
+            )}
+          >
+            Creative: {creative ? 'On' : 'Off'}
+          </button>
+          <span className="text-[11px] text-gray-500">
+            {creative
+              ? 'Expands the idea — places the subject in a scene and adds fitting details.'
+              : 'Faithful — structures the prompt as given, with a minimal background.'}
+          </span>
+        </div>
+
+        <div className="mb-3">
+          <label className="block text-[11px] mb-1 text-gray-400">Additional instructions (optional)</label>
+          <textarea
+            value={instructions}
+            onChange={e => setInstructions(e.target.value)}
+            disabled={isRunning || isDone}
+            rows={2}
+            placeholder="e.g. keep it a close-up portrait, prefer a daytime setting, always include a 9:16 vertical framing..."
+            className="w-full text-xs px-2 py-1.5 bg-gray-950 border border-gray-700 rounded-md text-gray-100 placeholder-gray-600 focus:ring-1 focus:ring-gray-600 focus:outline-none resize-none disabled:opacity-50"
+          />
+        </div>
+
+        <div className="rounded-md bg-gray-950 border border-gray-800 max-h-[45vh] overflow-y-auto divide-y divide-gray-800">
+          {prompts.map(p => {
+            const s = status[p.index];
+            const isSelected = !!selected[p.index];
+            const display = upsampled[p.index] ?? p.prompt;
+            const dim = (isRunning || isDone) && !s;
+            return (
+              <div key={p.index} className={classNames('flex items-start gap-2 px-3 py-2', { 'opacity-40': dim })}>
+                <input
+                  type="checkbox"
+                  checked={isSelected}
+                  disabled={isRunning || isDone}
+                  onChange={() => toggle(p.index)}
+                  className="mt-0.5 flex-shrink-0 accent-purple-600"
+                />
+                <div className="flex-shrink-0 mt-0.5">
+                  <StatusIcon s={s} />
+                </div>
+                <div className="min-w-0 flex-1">
+                  <div className="flex items-center gap-2">
+                    <span className="text-[10px] text-gray-500 flex-shrink-0">#{p.index + 1}</span>
+                    <span className="text-[10px] text-gray-500 flex-shrink-0">{p.aspectRatio}</span>
+                    {upsampled[p.index] && (
+                      <span className="text-[10px] text-emerald-400 flex-shrink-0">upsampled</span>
+                    )}
+                  </div>
+                  <div
+                    className={classNames('text-xs break-words line-clamp-2', {
+                      'text-emerald-200': upsampled[p.index],
+                      'text-gray-300': !upsampled[p.index],
+                    })}
+                    title={upsampled[p.index] ? upsampled[p.index] : p.prompt}
+                  >
+                    {display}
+                  </div>
+                </div>
+              </div>
+            );
+          })}
+        </div>
+
+        {(isRunning || isDone) && logOutput && (
+          <div
+            ref={logRef}
+            className="mt-3 font-mono text-[10px] whitespace-pre-wrap break-all overflow-y-auto rounded-md p-2 max-h-32 bg-black text-gray-300"
+          >
+            {logOutput}
+          </div>
+        )}
+
+        <div className="mt-4 flex justify-end gap-2">
+          <button
+            type="button"
+            onClick={onClose}
+            disabled={isRunning}
+            className="px-4 py-2 text-sm text-gray-300 hover:text-gray-100 disabled:opacity-40 disabled:cursor-not-allowed rounded-md"
+          >
+            {isDone ? 'Close' : 'Cancel'}
+          </button>
+          {!isDone && (
+            <button
+              type="button"
+              onClick={onRun}
+              disabled={isRunning || selectedCount === 0}
+              className="px-4 py-2 text-sm bg-purple-600 hover:bg-purple-700 disabled:opacity-40 disabled:cursor-not-allowed text-white rounded-md"
+            >
+              {isRunning ? 'Upsampling...' : `Upsample (${selectedCount})`}
+            </button>
+          )}
+        </div>
+      </div>
+    </Modal>
+  );
+};
+
+export default UpsamplePromptsModal;

ui_scripts/upsample_ideogram4_caption.py

@@ -0,0 +1,389 @@
+"""Upsample a short user idea into a full Ideogram4 structured-JSON caption.
+
+Runs the Ideogram4 generation ("magic prompt") system prompt through
+Qwen/Qwen3-VL-8B-Instruct as a text-only request and returns the resulting JSON.
+Nothing is written to disk -- the upsampled JSON object is printed to stdout
+(progress/logs go to stderr so stdout stays clean for the caller to parse).
+"""
+
+import argparse
+import json
+import os
+import re
+import sys
+from typing import Optional
+
+import torch
+
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+# Make the repo importable (e.g. `toolkit.util.quantize`) regardless of cwd.
+if REPO_ROOT not in sys.path:
+    sys.path.insert(0, REPO_ROOT)
+
+# The generation prompt lives here. It's a `name = """<content>"""` file, but the
+# content intentionally contains literal `\uNNNN` and `\n` sequences that are not
+# valid Python escapes, so it cannot be imported -- we read the triple-quoted
+# content verbatim as text instead.
+_PROMPT_PATH = os.path.join(
+    REPO_ROOT,
+    "extensions_built_in",
+    "captioner",
+    "prompts",
+    "ideogram4_upsample_prompt.py",
+)
+
+# Swapped into the prompt's {{mode_directive}} slot. Both keep the FIDELITY rules;
+# they only differ on how much the model may expand beyond the literal prompt.
+FAITHFUL_DIRECTIVE = (
+    "- **Fill in ONLY what the structure needs.** Add a concrete background shell, "
+    "bounding boxes, and the required elements/text -- nothing else. Do NOT add new "
+    "subjects, props, narrative, mood, or a setting the user did not specify. If the "
+    "prompt names no location, keep the background minimal. If the prompt is sparse, "
+    "the scene stays sparse."
+)
+
+CREATIVE_DIRECTIVE = (
+    "- **Expand the scene while keeping the user's idea intact.** Place the subject in "
+    "a specific, believable setting and build a real background environment with fitting "
+    "secondary details (props, depth layers, atmosphere) that serve the idea -- never a "
+    "blank or 'plain' background when a setting can be implied. Everything you add must "
+    "support, never replace or contradict, what the user asked for, and you must not "
+    "introduce a different main subject. The FIDELITY rules above still hold: triggers "
+    "verbatim, no invented appearance for a named person, no elaboration of a named style."
+)
+
+DTYPE_MAP = {
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "float16": torch.float16,
+    "fp16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+}
+
+
+def log(message: str) -> None:
+    print(message, file=sys.stderr, flush=True)
+
+
+def load_generation_prompt() -> str:
+    with open(_PROMPT_PATH, "r", encoding="utf-8") as f:
+        src = f.read()
+    # Extract the triple-quoted body verbatim (see note on _PROMPT_PATH).
+    start = src.find('"""')
+    end = src.rfind('"""')
+    if start == -1 or end <= start:
+        raise RuntimeError(f"Could not parse prompt body from {_PROMPT_PATH}")
+    return src[start + 3 : end]
+
+
+def build_prompt(
+    template: str,
+    aspect_ratio: str,
+    original_prompt: str,
+    creative: bool = False,
+    instructions: str = "",
+) -> str:
+    directive = CREATIVE_DIRECTIVE if creative else FAITHFUL_DIRECTIVE
+    prompt = template.replace("{{mode_directive}}", directive)
+    prompt = prompt.replace("{{user_instructions}}", instructions.strip() or "None.")
+    prompt = prompt.replace("{{aspect_ratio}}", aspect_ratio)
+    prompt = prompt.replace("{{original_prompt}}", original_prompt)
+    return prompt
+
+
+def extract_json(raw: str):
+    """Pull the JSON object out of the model output, tolerating code fences and
+    stray preamble. Returns the parsed dict or None."""
+    text = raw.strip()
+    fence = re.search(r"```(?:json)?\s*(.*?)```", text, re.DOTALL)
+    if fence:
+        text = fence.group(1).strip()
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end <= start:
+        return None
+    try:
+        return json.loads(text[start : end + 1])
+    except json.JSONDecodeError:
+        return None
+
+
+def sanitize_bbox(bbox):
+    """The generation prompt already emits normalized 0-1000 [y1,x1,y2,x2]. Clamp
+    to range, sort each axis pair, coerce to ints (keeps y/x order). Returns the
+    cleaned box or None to drop it."""
+    if not isinstance(bbox, (list, tuple)) or len(bbox) != 4:
+        return None
+    try:
+        y1, x1, y2, x2 = [float(v) for v in bbox]
+    except (TypeError, ValueError):
+        return None
+    y1, y2 = sorted((max(0, min(1000, round(y1))), max(0, min(1000, round(y2)))))
+    x1, x2 = sorted((max(0, min(1000, round(x1))), max(0, min(1000, round(x2)))))
+    if y2 <= y1 or x2 <= x1:
+        return None
+    return [y1, x1, y2, x2]
+
+
+def sanitize_caption(data: dict, aspect_ratio: str) -> dict:
+    """Light cleanup: force a concrete aspect ratio (never echo 'auto') and clean
+    each bbox. Leaves prose untouched."""
+    if aspect_ratio.lower() != "auto":
+        data["aspect_ratio"] = aspect_ratio
+    decon = data.get("compositional_deconstruction", {})
+    elements = decon.get("elements", [])
+    if isinstance(elements, list):
+        for el in elements:
+            if isinstance(el, dict) and "bbox" in el:
+                cleaned = sanitize_bbox(el["bbox"])
+                if cleaned is None:
+                    el.pop("bbox", None)
+                else:
+                    el["bbox"] = cleaned
+    return data
+
+
+def upsample_one(
+    model,
+    processor,
+    device,
+    template,
+    idea,
+    aspect_ratio,
+    gen_kwargs,
+    creative=False,
+    instructions="",
+) -> Optional[dict]:
+    """Run one idea through the generation prompt. Returns the cleaned caption
+    dict, or None if the model output couldn't be parsed."""
+    full_prompt = build_prompt(
+        template, aspect_ratio, idea.strip(), creative, instructions
+    )
+    messages = [{"role": "user", "content": [{"type": "text", "text": full_prompt}]}]
+    inputs = processor.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(device)
+
+    generated_ids = model.generate(**inputs, **gen_kwargs)
+    trimmed = [
+        out_ids[len(in_ids) :]
+        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0].strip()
+
+    data = extract_json(output_text)
+    if data is None:
+        log("Failed to parse JSON from model output. Raw output follows:")
+        log(output_text)
+        return None
+    return sanitize_caption(data, aspect_ratio)
+
+
+def normalize_item(item, default_aspect_ratio):
+    """Accept either a bare prompt string or {'prompt': ..., 'aspect_ratio': ...}.
+    Returns (idea, aspect_ratio) or None if the item is malformed/empty."""
+    if isinstance(item, str):
+        idea, aspect_ratio = item, default_aspect_ratio
+    elif isinstance(item, dict) and isinstance(item.get("prompt"), str):
+        idea = item["prompt"]
+        aspect_ratio = item.get("aspect_ratio") or default_aspect_ratio
+    else:
+        return None
+    if not idea.strip():
+        return None
+    return idea, aspect_ratio
+
+
+def load_model(
+    model_name_or_path: str,
+    dtype: torch.dtype,
+    device: torch.device,
+    quantize: bool,
+    qtype: str,
+):
+    from transformers import (
+        Qwen3VLForConditionalGeneration,
+        Qwen3VLMoeForConditionalGeneration,
+        AutoProcessor,
+    )
+
+    ModelClass = (
+        Qwen3VLMoeForConditionalGeneration
+        if "B-A" in model_name_or_path
+        else Qwen3VLForConditionalGeneration
+    )
+    log(f"Loading {model_name_or_path}")
+    model = ModelClass.from_pretrained(
+        model_name_or_path, dtype=dtype, device_map="cpu"
+    )
+    if quantize:
+        # Lazy import so the common (non-quantized) path needs no toolkit deps.
+        from optimum.quanto import freeze
+        from toolkit.util.quantize import quantize as quantize_model, get_qtype
+
+        log(f"Quantizing model ({qtype})")
+        quantize_model(model, weights=get_qtype(qtype))
+        freeze(model)
+    model.to(device)
+    model.eval()
+    processor = AutoProcessor.from_pretrained(model_name_or_path)
+    return model, processor
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Upsample a short idea into an Ideogram4 structured-JSON caption."
+    )
+    parser.add_argument(
+        "--prompt",
+        default=None,
+        help="A single user idea to upsample (prints one JSON object).",
+    )
+    parser.add_argument(
+        "--prompts",
+        default=None,
+        help=(
+            "JSON list to upsample in one model load (prints a JSON list, same order). "
+            'Each item is a prompt string or {"prompt": "...", "aspect_ratio": "W:H"}. '
+            "Failed/empty items come back as null."
+        ),
+    )
+    parser.add_argument(
+        "--aspect_ratio",
+        default="auto",
+        help="Default aspect ratio as 'W:H', or 'auto'. Per-item values override it.",
+    )
+    parser.add_argument("--model_name_or_path", default="Qwen/Qwen3-VL-8B-Instruct")
+    parser.add_argument("--max_new_tokens", type=int, default=3072)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--dtype", default="bf16", choices=list(DTYPE_MAP.keys()))
+    parser.add_argument("--quantize", action="store_true")
+    parser.add_argument("--qtype", default="float8")
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.7,
+        help="Sampling temperature. <= 0 uses greedy decoding.",
+    )
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument(
+        "--creative",
+        action="store_true",
+        help="Expand the prompt into a populated scene (default: faithful/minimal).",
+    )
+    parser.add_argument(
+        "--instructions",
+        default="",
+        help="Extra user instructions injected into the system prompt for every item.",
+    )
+    parser.add_argument("--pretty", action="store_true", help="Indent the output JSON.")
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        help=(
+            "Emit one compact JSON line per prompt as it completes "
+            '({"index": i, "caption": {...}|null}) instead of a single final list.'
+        ),
+    )
+    args = parser.parse_args()
+
+    if bool(args.prompt) == bool(args.prompts):
+        print(
+            "Provide exactly one of --prompt or --prompts.", file=sys.stderr, flush=True
+        )
+        return 2
+
+    # Resolve the work list up front so we can fail fast on bad input.
+    if args.prompts is not None:
+        try:
+            raw_items = json.loads(args.prompts)
+        except json.JSONDecodeError as e:
+            print(f"Failed to parse --prompts JSON: {e}", file=sys.stderr, flush=True)
+            return 2
+        if not isinstance(raw_items, list) or len(raw_items) == 0:
+            print(
+                "--prompts must be a non-empty JSON list.", file=sys.stderr, flush=True
+            )
+            return 2
+        batch = True
+    else:
+        if not args.prompt.strip():
+            print("--prompt must not be empty.", file=sys.stderr, flush=True)
+            return 2
+        raw_items = [args.prompt]
+        batch = False
+
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+
+    device = torch.device(args.device)
+    dtype = DTYPE_MAP[args.dtype]
+    indent = 2 if args.pretty else None
+
+    template = load_generation_prompt()
+
+    gen_kwargs = {"max_new_tokens": args.max_new_tokens}
+    if args.temperature and args.temperature > 0:
+        gen_kwargs.update(do_sample=True, temperature=args.temperature)
+    else:
+        gen_kwargs.update(do_sample=False)
+
+    with torch.no_grad():
+        model, processor = load_model(
+            args.model_name_or_path, dtype, device, args.quantize, args.qtype
+        )
+
+        results = []
+        for idx, item in enumerate(raw_items):
+            norm = normalize_item(item, args.aspect_ratio)
+            if norm is None:
+                log(f"[{idx + 1}/{len(raw_items)}] invalid/empty item, skipping")
+                result = None
+            else:
+                idea, aspect_ratio = norm
+                log(
+                    f"[{idx + 1}/{len(raw_items)}] Generating (aspect_ratio={aspect_ratio})..."
+                )
+                result = upsample_one(
+                    model,
+                    processor,
+                    device,
+                    template,
+                    idea,
+                    aspect_ratio,
+                    gen_kwargs,
+                    args.creative,
+                    args.instructions,
+                )
+            results.append(result)
+            # Stream each result on its own compact line so callers can update live.
+            if args.stream:
+                print(
+                    json.dumps({"index": idx, "caption": result}, ensure_ascii=False),
+                    flush=True,
+                )
+
+    if args.stream:
+        return 0 if any(r is not None for r in results) else 1
+
+    if batch:
+        print(json.dumps(results, ensure_ascii=False, indent=indent), flush=True)
+        # Non-zero only if nothing succeeded.
+        return 0 if any(r is not None for r in results) else 1
+
+    if results[0] is None:
+        return 1
+    print(json.dumps(results[0], ensure_ascii=False, indent=indent), flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())