add more control

app.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import random
 import time
 import traceback
 from io import BytesIO
@@ -28,6 +29,11 @@ MAX_RETRY_COUNT = int(os.environ.get("MAX_RETRY_COUNT", 3))
 POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", 2.0))
 MAX_POLL_TIME = int(os.environ.get("MAX_POLL_TIME", 300))
 
+# Local-only test mode: skip the real API and return a random-colored image after
+# a randomised delay. Useful for testing the UI / queueing flow without burning
+# real model credits. Enable with FAKE_TEST=1 (also accepts "true"/"yes").
+FAKE_TEST = os.environ.get("FAKE_TEST", "").strip().lower() in ("1", "true", "yes", "on")
+
 # Predefined aspect ratios (wh_ratio) — kept in the same order as the original
 # (width, height) list so each entry mirrors the previous resolution choice:
 #   1:1  ←→ 2048×2048   4:3  ←→ 2304×1728   3:4  ←→ 1728×2304
@@ -54,6 +60,8 @@ logger.info(
 logger.info(
     f"Retry configuration: MAX_RETRY_COUNT={MAX_RETRY_COUNT}, POLL_INTERVAL={POLL_INTERVAL}s, MAX_POLL_TIME={MAX_POLL_TIME}s"
 )
+if FAKE_TEST:
+    logger.warning("FAKE_TEST mode is ENABLED — no real API calls will be made.")
 
 
 class APIError(Exception):
@@ -81,13 +89,27 @@ def _headers() -> dict:
     return {"Authorization": f"Bearer {API_TOKEN}"}
 
 
-def create_request(prompt, wh_ratio):
+def create_request(
+    prompt,
+    wh_ratio,
+    negative_prompt="",
+    enable_prompt_refine=True,
+    seed=-1,
+    guidance_scale=5.0,
+):
     """
     Submit an image generation request to the API.
 
     Args:
         prompt (str): Text prompt describing the image to generate
         wh_ratio (str): Aspect ratio for the output image (e.g. "16:9")
+        negative_prompt (str): Optional text describing what to avoid.
+        enable_prompt_refine (bool): Whether to let the backend rewrite/expand
+            the prompt before generation. Sent to the API as 0 / 1.
+        seed (int): Generation seed. -1 means the backend will pick one
+            randomly; any other integer fixes the seed for reproducible runs.
+        guidance_scale (float): Classifier-free guidance strength. Higher
+            values follow the prompt more strictly.
 
     Returns:
         str: Task ID
@@ -96,7 +118,10 @@ def create_request(prompt, wh_ratio):
         APIError: If the API request fails
     """
     logger.info(
-        f"Starting create_request with prompt='{prompt[:50]}...', wh_ratio={wh_ratio}"
+        f"Starting create_request with prompt='{prompt[:50]}...', "
+        f"wh_ratio={wh_ratio}, enable_prompt_refine={enable_prompt_refine}, "
+        f"seed={seed}, guidance_scale={guidance_scale}, "
+        f"negative_prompt='{(negative_prompt or '')[:30]}...'"
     )
 
     if not prompt or not prompt.strip():
@@ -107,11 +132,27 @@ def create_request(prompt, wh_ratio):
         logger.error(f"Invalid wh_ratio: {wh_ratio}. Valid options: {WH_RATIO_OPTIONS}")
         raise ValueError(f"Invalid aspect ratio. Must be one of: {', '.join(WH_RATIO_OPTIONS)}")
 
+    try:
+        seed_int = int(seed)
+    except (TypeError, ValueError):
+        logger.warning(f"Invalid seed value '{seed}', falling back to -1 (random)")
+        seed_int = -1
+
+    try:
+        guidance_scale_f = float(guidance_scale)
+    except (TypeError, ValueError):
+        logger.warning(f"Invalid guidance_scale '{guidance_scale}', falling back to 5.0")
+        guidance_scale_f = 5.0
+
     model_params = {
         "prompt": prompt,
         "wh_ratio": wh_ratio,
         "model_id": MODEL_ID,
         "n": 1,
+        "negative_prompt": negative_prompt or "",
+        "enable_prompt_refine": 1 if enable_prompt_refine else 0,
+        "seed": seed_int,
+        "guidance_scale": guidance_scale_f,
     }
 
     url = _build_request_url()
@@ -519,6 +560,177 @@ textarea:focus, input:focus,
     background: #ffffff !important;
 }
 
+/* Negative prompt — softer accent so it visually de-emphasises vs. the main prompt */
+.negative-prompt textarea {
+    background: #fbfbfd !important;
+    border-color: #e3e3e8 !important;
+}
+.negative-prompt textarea:focus {
+    background: #ffffff !important;
+}
+
+/* Advanced options row — keeps the refine switch + seed input visually paired */
+.advanced-row {
+    gap: 14px !important;
+    margin-top: 2px;
+}
+
+/* Refine toggle — iOS-style settings card.
+   Goal: title + helper text on the left, a polished pill switch on the right,
+   everything contained inside a single soft card so the helper text no longer
+   floats orphaned above the box. */
+.refine-toggle {
+    background: linear-gradient(180deg, #ffffff 0%, #f5f5f7 100%) !important;
+    border-radius: 14px !important;
+    border: 1px solid rgba(0,0,0,0.06) !important;
+    padding: 14px 16px !important;
+    box-shadow: 0 1px 2px rgba(0,0,0,0.03) !important;
+    transition: border-color 0.18s ease, box-shadow 0.18s ease !important;
+    min-height: 88px;
+    display: flex !important;
+    flex-direction: column !important;
+    justify-content: center !important;
+}
+.refine-toggle:hover {
+    border-color: rgba(0,0,0,0.10) !important;
+    box-shadow: 0 1px 2px rgba(0,0,0,0.04), 0 4px 14px rgba(0,0,0,0.05) !important;
+}
+
+/* Strip default backgrounds from gradio's inner wrappers so only our card shows. */
+.refine-toggle .form,
+.refine-toggle .wrap,
+.refine-toggle .form-wrap,
+.refine-toggle > div {
+    background: transparent !important;
+    border: none !important;
+    padding: 0 !important;
+    margin: 0 !important;
+    box-shadow: none !important;
+}
+
+/* Helper / "info" text becomes a proper subtitle UNDER the toggle row. */
+.refine-toggle [data-testid="block-info"],
+.refine-toggle .info {
+    color: #6e6e73 !important;
+    font-size: 12px !important;
+    line-height: 1.4 !important;
+    margin: 8px 0 0 0 !important;
+    padding: 0 !important;
+    text-align: left !important;
+    order: 2 !important;
+}
+
+/* Force the gradio wrapper to stack: label-row first, info below. */
+.refine-toggle .form,
+.refine-toggle > div:not([data-testid="block-info"]):not(.info) {
+    display: flex !important;
+    flex-direction: column !important;
+    align-items: stretch !important;
+}
+
+/* Label row: title on the LEFT (full width), toggle pinned to the RIGHT. */
+.refine-toggle label {
+    display: flex !important;
+    align-items: center !important;
+    justify-content: space-between !important;
+    flex-direction: row-reverse !important;
+    cursor: pointer !important;
+    margin: 0 !important;
+    padding: 0 !important;
+    gap: 14px !important;
+    width: 100% !important;
+    order: 1 !important;
+}
+.refine-toggle label > span {
+    color: #1d1d1f !important;
+    font-size: 15px !important;
+    font-weight: 600 !important;
+    letter-spacing: -0.01em;
+    flex: 1 1 auto !important;
+    text-align: left !important;
+    line-height: 1.3 !important;
+}
+
+/* Pill switch — bigger, smoother, more "Apple-like" */
+.refine-toggle input[type="checkbox"] {
+    appearance: none;
+    -webkit-appearance: none;
+    width: 46px !important;
+    height: 28px !important;
+    border-radius: 999px !important;
+    background: #e5e5ea !important;
+    position: relative;
+    cursor: pointer;
+    transition: background 0.22s ease, box-shadow 0.22s ease;
+    border: none !important;
+    flex-shrink: 0 !important;
+    margin: 0 !important;
+    box-shadow: inset 0 0 1px rgba(0,0,0,0.06);
+}
+.refine-toggle input[type="checkbox"]::after {
+    content: "";
+    position: absolute;
+    top: 2px;
+    left: 2px;
+    width: 24px;
+    height: 24px;
+    border-radius: 50%;
+    background: #ffffff;
+    box-shadow: 0 2px 5px rgba(0,0,0,0.18), 0 0 1px rgba(0,0,0,0.05);
+    transition: transform 0.24s cubic-bezier(0.4, 0.0, 0.2, 1);
+}
+.refine-toggle input[type="checkbox"]:hover {
+    background: #dcdce0 !important;
+}
+.refine-toggle input[type="checkbox"]:checked {
+    background: #34c759 !important;
+}
+.refine-toggle input[type="checkbox"]:checked:hover {
+    background: #30b352 !important;
+}
+.refine-toggle input[type="checkbox"]:checked::after {
+    transform: translateX(18px);
+}
+.refine-toggle input[type="checkbox"]:focus-visible {
+    box-shadow: 0 0 0 4px rgba(0,113,227,0.20) !important;
+}
+.refine-toggle input[type="checkbox"]:active::after {
+    /* tiny squish on press, very iOS */
+    width: 28px;
+}
+.refine-toggle input[type="checkbox"]:checked:active::after {
+    transform: translateX(14px);
+}
+
+/* Seed number input — match the prompt/dropdown rounding */
+.seed-input input[type="number"] {
+    border-radius: 12px !important;
+    padding: 10px 14px !important;
+    font-variant-numeric: tabular-nums;
+}
+/* Hide the native spinner buttons on number inputs for a cleaner look */
+.seed-input input[type="number"]::-webkit-outer-spin-button,
+.seed-input input[type="number"]::-webkit-inner-spin-button {
+    -webkit-appearance: none;
+    margin: 0;
+}
+.seed-input input[type="number"] {
+    -moz-appearance: textfield;
+}
+/* Keep the seed column visually aligned with the refine card next to it */
+.seed-input {
+    align-self: stretch !important;
+}
+
+/* Guidance scale slider — Apple-blue track + softer thumb */
+.guidance-slider input[type="range"] {
+    accent-color: #0071e3 !important;
+}
+.guidance-slider .head { padding-top: 0 !important; }
+.guidance-slider {
+    margin-top: 4px;
+}
+
 /* Footer tagline */
 .tagline {
     text-align: center;
@@ -603,6 +815,144 @@ def _status_html(text: str, kind: str = "info") -> str:
     )
 
 
+def _queue_text(waiting: int, running: int) -> str:
+    """Inline status text describing the queue from the user's POV.
+
+    Used inside the right-side status pill while THIS user's request is
+    sitting in the queue waiting for an open slot.
+    """
+    if waiting <= 0 and running <= 0:
+        return "Queued · waiting for an open slot…"
+    parts = []
+    if waiting > 0:
+        parts.append(f"{waiting} waiting")
+    if running > 0:
+        parts.append(f"{running} generating")
+    return f"In queue · {' · '.join(parts)}"
+
+
+def _read_queue_stats(demo_obj) -> tuple[int, int, float]:
+    """Best-effort read of gradio's internal queue.
+
+    Gradio 5.x structure (gradio/queueing.py):
+      demo._queue.event_queue_per_concurrency_id: dict[str, EventQueue]
+        EventQueue.queue: list[Event]              ← actually-waiting events
+      demo._queue.active_jobs: list[None | list[Event]]
+        each slot is None (idle) or a list of currently-processing events.
+      demo._queue.process_time_per_fn: dict[BlockFunction, ProcessTime]
+        ProcessTime.avg_time: float
+
+    Returns (waiting, running, avg_secs). Each lookup is wrapped in a try
+    so the UI degrades gracefully ("idle") if Gradio ever renames a field.
+    """
+    try:
+        q = getattr(demo_obj, "_queue", None)
+        if q is None:
+            return 0, 0, 0.0
+
+        # ---- Waiting: events sitting in EventQueue.queue ----
+        waiting = 0
+        events_per_cid = getattr(q, "event_queue_per_concurrency_id", None) or {}
+        for ev_q in events_per_cid.values():
+            # Newer gradio: ev_q is an EventQueue with a .queue list.
+            # Older/alt: ev_q might already be a list. Handle both.
+            inner = getattr(ev_q, "queue", ev_q)
+            try:
+                waiting += len(inner)
+            except (TypeError, AttributeError):
+                # Last resort: try iterating
+                try:
+                    waiting += sum(1 for _ in inner)
+                except Exception:
+                    continue
+
+        # ---- Running: count events held in active_jobs slots ----
+        # Each slot is None or a list[Event]; sum the list lengths.
+        running = 0
+        active = getattr(q, "active_jobs", None) or []
+        for slot in active:
+            if slot is None:
+                continue
+            try:
+                running += len(slot)
+            except (TypeError, AttributeError):
+                running += 1  # very old gradio: single Event per slot
+
+        # ---- Average per-run time (best effort across versions) ----
+        avg_secs = 0.0
+        # Gradio 5.x: dict[BlockFunction, ProcessTime] with .avg_time
+        ptpf = getattr(q, "process_time_per_fn", None)
+        if isinstance(ptpf, dict) and ptpf:
+            try:
+                vals = []
+                for v in ptpf.values():
+                    avg_t = getattr(v, "avg_time", None)
+                    if avg_t:
+                        vals.append(float(avg_t))
+                if vals:
+                    avg_secs = sum(vals) / len(vals)
+            except Exception:
+                avg_secs = 0.0
+        # Older: dict[int, float]
+        if not avg_secs:
+            ptpfi = getattr(q, "process_time_per_fn_index", None)
+            if isinstance(ptpfi, dict) and ptpfi:
+                try:
+                    vals = [float(v) for v in ptpfi.values() if v]
+                    if vals:
+                        avg_secs = sum(vals) / len(vals)
+                except Exception:
+                    avg_secs = 0.0
+        # Oldest: single float on the queue
+        if not avg_secs:
+            apt = getattr(q, "avg_process_time", None)
+            if apt:
+                try:
+                    avg_secs = float(apt)
+                except Exception:
+                    avg_secs = 0.0
+
+        return waiting, running, avg_secs
+    except Exception as exc:
+        logger.debug(f"Queue introspection failed: {exc}")
+        return 0, 0, 0.0
+
+
+def _fake_generation_iter(prompt: str, wh_ratio_value: str):
+    """FAKE_TEST mode generator.
+
+    Mimics the real flow's yield protocol without hitting any external API.
+    Useful for exercising the queue UI / status pill locally.
+
+    Yields (image_or_None, status_html) tuples. The very first 'Sending
+    request to API…' yield is emitted by the caller, so this iterator picks
+    up from 'Request submitted' onwards.
+    """
+    time.sleep(random.uniform(0.4, 1.0))
+
+    fake_id = f"{random.randint(0, 0xFFFFFFFF):08x}"
+    yield None, _status_html(f"Request submitted · Task {fake_id}…", "running")
+
+    target_secs = random.uniform(8.0, 22.0)
+    start = time.time()
+    while time.time() - start < target_secs:
+        elapsed = int(time.time() - start)
+        yield None, _status_html(f"Generating… {elapsed}s", "running")
+        time.sleep(POLL_INTERVAL)
+
+    yield None, _status_html("Downloading image…", "running")
+    time.sleep(0.3)
+
+    rgb = (
+        random.randint(40, 220),
+        random.randint(40, 220),
+        random.randint(40, 220),
+    )
+    fake_image = Image.new("RGB", (1024, 1024), color=rgb)
+    logger.info(f"FAKE_TEST: returning random color image rgb={rgb}, took {target_secs:.1f}s")
+    yield fake_image, _status_html("Image generated", "success")
+
+
 def create_ui():
     logger.info("Creating Gradio UI")
     with gr.Blocks(
@@ -610,6 +960,14 @@ def create_ui():
         theme=APPLE_THEME,
         css=APPLE_CSS,
     ) as demo:
+        # Per-session state used to drive the right-side status pill while
+        # this user's request is sitting in the queue. `last_pill_state`
+        # caches the most recent pill HTML so the timer can return
+        # `gr.update()` (=no-op, no DOM replacement → no flicker) when the
+        # queue counts haven't actually changed between ticks.
+        queued_state = gr.State(False)
+        last_pill_state = gr.State("")
+
         with gr.Row(equal_height=False):
             with gr.Column(scale=1, elem_classes=["panel-card"]):
                 prompt = gr.Textbox(
@@ -619,6 +977,14 @@ def create_ui():
                     show_label=True,
                 )
 
+                negative_prompt = gr.Textbox(
+                    label="Negative Prompt",
+                    placeholder="Things you want to avoid in the image (optional)...",
+                    lines=2,
+                    show_label=True,
+                    elem_classes=["negative-prompt"],
+                )
+
                 wh_ratio = gr.Dropdown(
                     choices=WH_RATIO_OPTIONS,
                     value=WH_RATIO_OPTIONS[0],
@@ -626,6 +992,34 @@ def create_ui():
                     info="Width : Height",
                 )
 
+                guidance_scale = gr.Slider(
+                    minimum=1.0,
+                    maximum=20.0,
+                    step=0.1,
+                    value=5.0,
+                    label="Guidance Scale",
+                    info="Higher values follow the prompt more strictly",
+                    elem_classes=["guidance-slider"],
+                )
+
+                with gr.Row(elem_classes=["advanced-row"], equal_height=True):
+                    enable_prompt_refine = gr.Checkbox(
+                        value=True,
+                        label="Prompt Refine",
+                        info="Let the model rewrite & enrich your prompt",
+                        elem_classes=["refine-toggle"],
+                        scale=1,
+                    )
+                    seed = gr.Number(
+                        value=-1,
+                        label="Seed",
+                        info="Use -1 for a random seed",
+                        precision=0,
+                        minimum=-1,
+                        elem_classes=["seed-input"],
+                        scale=1,
+                    )
+
                 with gr.Row():
                     clear_btn = gr.Button("Clear", variant="secondary", scale=1)
                     generate_btn = gr.Button("Generate", variant="primary", scale=3)
@@ -645,9 +1039,20 @@ def create_ui():
                     elem_classes=["image-output"],
                 )
 
-        def generate_with_status(prompt, wh_ratio_value):
+        def generate_with_status(
+            prompt,
+            wh_ratio_value,
+            negative_prompt_value,
+            enable_prompt_refine_value,
+            seed_value,
+            guidance_scale_value,
+        ):
             logger.info(
-                f"Starting image generation with prompt='{(prompt or '')[:50]}...', wh_ratio={wh_ratio_value}"
+                f"Starting image generation with prompt='{(prompt or '')[:50]}...', "
+                f"wh_ratio={wh_ratio_value}, "
+                f"negative_prompt='{(negative_prompt_value or '')[:30]}...', "
+                f"enable_prompt_refine={enable_prompt_refine_value}, "
+                f"seed={seed_value}, guidance_scale={guidance_scale_value}"
             )
 
             yield None, _status_html("Sending request to API…", "running")
@@ -663,8 +1068,30 @@ def create_ui():
                     yield None, _status_html(f"Invalid aspect ratio “{wh_ratio_value}”", "error")
                     return
 
+                try:
+                    seed_int = int(seed_value) if seed_value is not None else -1
+                except (TypeError, ValueError):
+                    seed_int = -1
+
+                try:
+                    guidance_scale_f = float(guidance_scale_value) if guidance_scale_value is not None else 5.0
+                except (TypeError, ValueError):
+                    guidance_scale_f = 5.0
+
+                if FAKE_TEST:
+                    logger.info("FAKE_TEST mode active — bypassing real API call")
+                    yield from _fake_generation_iter(prompt, wh_ratio_value)
+                    return
+
                 logger.info("Creating API request")
-                task_id = create_request(prompt, wh_ratio_value)
+                task_id = create_request(
+                    prompt,
+                    wh_ratio_value,
+                    negative_prompt=negative_prompt_value or "",
+                    enable_prompt_refine=bool(enable_prompt_refine_value),
+                    seed=seed_int,
+                    guidance_scale=guidance_scale_f,
+                )
                 yield None, _status_html(f"Request submitted · Task {task_id[:8]}…", "running")
 
                 start_time = time.time()
@@ -752,21 +1179,90 @@ def create_ui():
                 logger.error(f"Full traceback: {traceback.format_exc()}")
                 yield None, _status_html(f"Unexpected error: {str(e)}", "error")
 
+        def _enter_queue():
+            """Click handler #1 — runs IMMEDIATELY (queue=False).
+
+            Flips queued_state to True and seeds the pill with a snapshot of
+            the current queue. Subsequent live updates come from the timer.
+            """
+            waiting, running, _ = _read_queue_stats(demo)
+            html = _status_html(_queue_text(waiting, running), "running")
+            return None, html, True, html
+
+        def _generate_wrapped(
+            prompt_value,
+            wh_ratio_value,
+            negative_prompt_value,
+            enable_prompt_refine_value,
+            seed_value,
+            guidance_scale_value,
+        ):
+            """Click handler #2 — queued.
+
+            Wraps the existing generator and, on the FIRST yield, also flips
+            queued_state to False so the timer stops touching the pill and
+            lets the generator's own `yield`s drive it (Generating XXs → ...).
+            """
+            first = True
+            for image, status_html in generate_with_status(
+                prompt_value,
+                wh_ratio_value,
+                negative_prompt_value,
+                enable_prompt_refine_value,
+                seed_value,
+                guidance_scale_value,
+            ):
+                if first:
+                    first = False
+                    yield image, status_html, False
+                else:
+                    yield image, status_html, gr.update()
+
         generate_btn.click(
-            fn=generate_with_status,
-            inputs=[prompt, wh_ratio],
-            outputs=[output_image, status_msg],
+            fn=_enter_queue,
+            inputs=None,
+            outputs=[output_image, status_msg, queued_state, last_pill_state],
+            queue=False,
             show_progress="hidden",
+        ).then(
+            fn=_generate_wrapped,
+            inputs=[prompt, wh_ratio, negative_prompt, enable_prompt_refine, seed, guidance_scale],
+            outputs=[output_image, status_msg, queued_state],
+            show_progress="minimal",
+            show_progress_on=[generate_btn],
         )
 
         def clear_outputs():
             logger.info("Clearing UI outputs")
-            return None, _status_html("Ready", "info")
+            return None, _status_html("Ready", "info"), False, ""
 
         clear_btn.click(
             fn=clear_outputs,
             inputs=None,
-            outputs=[output_image, status_msg],
+            outputs=[output_image, status_msg, queued_state, last_pill_state],
+        )
+
+        # Live queue updates inside the right-side pill — only while THIS
+        # user is queued. Returns gr.update() (no-op) when nothing changed,
+        # which prevents the DOM from being replaced and avoids the pulse
+        # animation resetting (= no flicker).
+        pill_timer = gr.Timer(value=1.5, active=True)
+
+        def _tick_pill(queued_flag, last_html):
+            if not queued_flag:
+                return gr.update(), last_html
+            waiting, running, _ = _read_queue_stats(demo)
+            new_html = _status_html(_queue_text(waiting, running), "running")
+            if new_html == last_html:
+                return gr.update(), last_html
+            return new_html, new_html
+
+        pill_timer.tick(
+            fn=_tick_pill,
+            inputs=[queued_state, last_pill_state],
+            outputs=[status_msg, last_pill_state],
+            queue=False,
+            show_progress="hidden",
         )
 
         gr.HTML(