Usage counter: append one metadata-only JSON line per generation to a private dataset (all replicas)

app.py

@@ -25,6 +25,7 @@ import io
 import json
 import random
 import re
+import socket
 import time
 from pathlib import Path
 
@@ -35,6 +36,7 @@ import compare
 from PIL import Image, ImageOps
 from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline, QwenImage21Transformer2DModel
 from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import calculate_dimensions
+from huggingface_hub import CommitScheduler
 
 # Students in the model repo. STUDENT="lora" (default): the v0.2.1 LoRA (r256) at the repo root, applied at runtime on
 # top of the base transformer and never merged (merging into bf16 keeps only ~47% of the adapter delta, the
@@ -47,6 +49,30 @@ STUDENT = os.environ.get("STUDENT", "lora")
 LORA_FILE = os.environ.get("LORA_FILE", "Qwen-Image-2.1-viggle-turbo-v0.2.1-6step-lora-r256.safetensors")
 STUDENT_TAG = {"full": "v0.1 full fine-tune (transformer/)", "lora": "v0.2.1 LoRA r256" if "v0.2.1" in LORA_FILE else f"LoRA {LORA_FILE}"}[STUDENT]
 HF_TOKEN = os.environ.get("HF_TOKEN")  # only needed while the model repo is private
+# Usage counter. The Space autoscales to several replicas and its log view shows only one of them, so the [gen] log lines
+# undercount. Every generate() call that reaches the GPU also appends one JSON line (metadata only: no prompt, no images)
+# to this container's own file, which CommitScheduler pushes to a private dataset every 10 minutes and at exit; one file
+# per container, so replicas never write the same path. The token is STATS_TOKEN, else HF_TOKEN; it needs write access to
+# that dataset. Without a token, or if the scheduler cannot start with it, the counter is off and the app runs as before.
+STATS_DIR = Path("/tmp/gen_stats")
+STATS_DIR.mkdir(exist_ok=True)
+STATS_FILE = STATS_DIR / f"{time.strftime('%Y%m%dT%H%M%SZ', time.gmtime())}-{socket.gethostname()}.jsonl"
+STATS_TOKEN = os.environ.get("STATS_TOKEN") or HF_TOKEN
+stats = None
+if STATS_TOKEN:
+    try:
+        stats = CommitScheduler(repo_id="Viggle/qwen-image-21-turbo-space-stats", repo_type="dataset", folder_path=STATS_DIR,
+                                path_in_repo="gens", every=10, private=True, token=STATS_TOKEN)
+    except Exception as e:  # e.g. a token that cannot see the dataset: count nothing rather than fail to start
+        print(f"[stats] off: {type(e).__name__}: {e}", flush=True)
+
+
+def record(**row):
+    if stats:
+        with stats.lock, STATS_FILE.open("a") as f:
+            f.write(json.dumps({"ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), **row}) + "\n")
+
+
 # The LoRA is sampled on 6 Euler steps whose raw (pre-shift) sigma nodes are RAW_NODES: the 4-step training schedule
 # linspace(1, 1/4, 4) with its highest-noise segment [1, 0.75] cut into three (1, 0.9375, 0.875, 0.75). The composition
 # is decided in that segment, and one big Euler step there ghosts and drifts the layout; the low-noise nodes 0.75, 0.5,
@@ -209,6 +235,7 @@ def denoise(prompt, images, width, height, steps, seed):
 
 
 def generate(prompt, refs, size_label, seed, randomize_seed, enhance="Auto", steps=STEPS, custom_width=1024, custom_height=1024):
+    t0 = time.perf_counter()
     steps = min(max(int(steps), 3), 8)  # 6 (RAW_NODES) is the validated schedule; 3-8 are accepted, fewer or more degrade quickly
     # refs: the gallery's (filepath, caption) pairs in upload order, which is the order the prompt's "image 1, 2, ..." refer
     # to. EXIF rotation + RGB is what the gr.Image slots used to do.
@@ -243,11 +270,19 @@ def generate(prompt, refs, size_label, seed, randomize_seed, enhance="Auto", ste
         used_prompt = enhance_prompt(prompt, images, ratio_name) or prompt
         pe_time = time.perf_counter() - started
         enhance_note = f" · enhanced {pe_time:.1f}s" if used_prompt != prompt else " · enhancement failed, prompt sent as written"
-    result, elapsed = denoise(used_prompt, images, width, height, steps, seed)
+    mode = "edit" if images else "t2i"
+    try:
+        result, elapsed = denoise(used_prompt, images, width, height, steps, seed)
+    except Exception as e:  # ZeroGPU quota / GPU timeout errors land here: count them, then show them as before
+        record(ok=False, mode=mode, refs=len(images), steps=steps, error=f"{type(e).__name__}: {e}"[:200],
+               total=round(time.perf_counter() - t0, 2))
+        raise
     # one line per generation, for counting usage from the Space logs (no prompt text)
-    print(f"[gen] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} mode={'edit' if images else 't2i'} refs={len(images)} "
+    print(f"[gen] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} mode={mode} refs={len(images)} "
           f"steps={steps} size={result.width}x{result.height} enhance={int(used_prompt != prompt)} pe={pe_time:.1f}s gen={elapsed:.2f}s",
           flush=True)
+    record(ok=True, mode=mode, refs=len(images), steps=steps, size=f"{result.width}x{result.height}", enhance=enhance,
+           enhanced=used_prompt != prompt, pe=round(pe_time, 2), gen=round(elapsed, 2), total=round(time.perf_counter() - t0, 2))
     return result, f"seed `{seed}` · {result.width}×{result.height} · {steps} steps · {elapsed:.2f}s · {STUDENT_TAG}{enhance_note}{clamped}", used_prompt