v0.2.1: xlarge card, no VAE tiling (tiled reference encode wrecked edits; 3-reference edit OOMed at 44.3 GiB on large)

README.md

@@ -22,8 +22,8 @@ tags:
   - dmd
 # Hardware: this Space needs ZeroGPU, set in Space Settings. `suggested_hardware` is deliberately
 # unset, because the only ZeroGPU value the Hub metadata accepts is the legacy `zero-a10g`, and a
-# 24 GB A10G cannot hold this pipeline. `app.py` requests the 48 GB card with
-# @spaces.GPU(size="large"); see the Hardware section for the measurements behind that.
+# 24 GB A10G cannot hold this pipeline. `app.py` requests the 96 GB card with
+# @spaces.GPU(size="xlarge") (no VAE tiling); see the Hardware section for the measurements behind that.
 # Secrets: set HF_TOKEN (read scope) while Viggle/Qwen-Image-2.1-viggle-turbo is private.
 # `hf_oauth` is not needed: the app calls no user-scoped Hub API.
 ---
@@ -148,41 +148,30 @@ convenience layered on top of the recipe, not part of it: untick it when the wor
 Weights are ~31.5 GiB resident in bf16 (13.25 transformer + 16.33 text encoder + 0.63 VAE + the r=256
 adapter, 1.3 GiB, loaded in bf16). Measured on one B200 with the v0.2 LoRA at 5 steps, prompt enhancement on,
 around a single `generate()` call including the prompt rewrite and the VAE encode/decode, with the allocator
-cache dropped before each measurement (`release/measure_mem.py`); the two 2048² rows are the v0.1 full-transformer
-measurements at 4 steps plus the adapter:
-
-| call | output | VAE tiling | peak allocated | peak reserved | wall clock (B200) |
-|---|---|---|---|---|---|
-| text-to-image | 1024×1024 | off | 38.24 GiB | 40.05 GiB | 0.8 s |
-| text-to-image | 1024×1024 | on | 32.39 GiB | 32.90 GiB | 1.1 s |
-| text-to-image | 2048×2048 | off | ~58.3 GiB | ~66 GiB | 3.4 s (v0.1) |
-| text-to-image | 2048×2048 | **on** | ~34.8 GiB | ~35.5 GiB | 4.7 s (v0.1) |
-| edit, 1 reference, auto | 1024×1024 | off | 40.13 GiB | 42.68 GiB | 1.1 s |
-| edit, 1 reference, auto | 1024×1024 | **on** | 34.91 GiB | 36.36 GiB | 1.8 s |
-| edit, 2 references, auto | 832×1248 | off | 42.04 GiB | 44.71 GiB | 1.5 s |
-| edit, 2 references, auto | 832×1248 | **on** | 37.52 GiB | 39.82 GiB | 2.3 s |
-| edit, 3 references, auto (first example) | 928×1152 | off | **44.29 GiB** | 47.84 GiB | 2.9 s |
-| edit, 3 references, auto (first example) | 928×1152 | **on** | 40.29 GiB | 43.58 GiB | 3.0 s |
-| edit, 3 references (largest) | 1760×1344 | off | 41.11 GiB | 45.14 GiB | 5.2 s |
-| edit, 3 references (largest) | 1760×1344 | **on** | **41.11 GiB** | 45.14 GiB | 4.8 s |
-
-The untiled VAE decode is the peak of every call at the 1024² bucket once references are resident, and
-of every larger text-to-image; the three-reference edit at 1760×1344 peaks in the denoising loop
-instead (41.1 GiB, the same tiled or not). ZeroGPU **`large`** is a 48 GB card with 44.7 GiB usable,
-so the untiled three-reference edit at 44.3 GiB is exactly the call that failed on the Space (the CUDA
-OOM surfaces there as `NVML_SUCCESS == r INTERNAL ASSERT FAILED` from the caching allocator, because the
-container cannot query NVML for the OOM report). `app.py` therefore requests `@spaces.GPU(duration=90,
-size="large")` at 1× quota and switches tiling on for every editing call and for the 1536²/2048²
-text-to-image buckets; 1024²-bucket text-to-image stays untiled and bit-identical to the untiled
-pipeline, tiled outputs differ from untiled ones by 1.6/255 mean absolute (measured at 2048×2048).
-With that, the largest call on the menu peaks at 41.1 GiB allocated, ~3.5 GiB inside the card; the
-*reserved* figures are what the caching allocator held on a card with room to spare - on the 44.7 GiB
-card it releases cached blocks before failing, so the allocated peak is the binding number, and
-`app.py` sets `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` to keep fragmentation down. Prompt
-enhancement runs before the denoise with nothing else live (~2 GiB of KV cache over the resident
-weights) and does not raise the peak. Wall-clock on the ZeroGPU card (half an RTX Pro 6000 Blackwell)
-will be slower than the B200 numbers above, and the prompt rewrite adds 2-8 s on a B200; `duration=90`
-covers both.
+cache dropped before each measurement (`release/measure_mem.py`; the two 2048² rows are the v0.1 full-transformer
+measurements at 4 steps plus the adapter). **No VAE tiling anywhere**: `vae.enable_tiling()` also tiles the
+*encode* of the reference images (256 px tiles), and tiled reference latents wreck an edit - duplicated subjects,
+wrong scale - while a tiled decode is not the validated pipeline either.
+
+| call | output | peak allocated | peak reserved | wall clock (B200) |
+|---|---|---|---|---|
+| text-to-image | 1024×1024 | 38.24 GiB | 40.05 GiB | 0.8 s |
+| text-to-image | 2048×2048 | ~58.3 GiB | ~66 GiB | 3.4 s (v0.1) |
+| edit, 1 reference, auto | 1024×1024 | 40.13 GiB | 42.68 GiB | 1.1 s |
+| edit, 2 references, auto | 832×1248 | 42.04 GiB | 44.71 GiB | 1.5 s |
+| edit, 3 references, auto (first example) | 928×1152 | 44.29 GiB | 47.84 GiB | 2.9 s |
+| edit, 3 references (largest) | 1760×1344 | ~52.8 GiB | ~59 GiB | 2.9 s (v0.1) |
+
+ZeroGPU `large` is a 48 GB card with 44.7 GiB usable, so the three-reference edit at the 1024² bucket is
+exactly the call that failed there (the CUDA OOM surfaces on ZeroGPU as `NVML_SUCCESS == r INTERNAL ASSERT
+FAILED` from the caching allocator, because the container cannot query NVML for the OOM report), and the
+2048² text-to-image and 1536² editing buckets never fit untiled. `app.py` therefore requests
+`@spaces.GPU(duration=90, size="xlarge")`: the full 96 GB RTX Pro 6000 Blackwell, at **2× quota** (an
+unauthenticated visitor's 2 min/day covers about three of the largest calls, a free account's 5 min about
+seven; PRO 40 min). `app.py` also sets `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. Prompt enhancement
+runs before the denoise with nothing else live (~2 GiB of KV cache over the resident weights) and does not
+raise the peak. Wall-clock on the ZeroGPU card will be slower than the B200 numbers above, and the prompt
+rewrite adds 2-8 s on a B200; `duration=90` covers both.
 
 **Cold start** is dominated by the download: ~30.9 GiB of base safetensors from
 `Qwen/Qwen-Image-2.1` (13.25 transformer + 16.33 text encoder + 1.26 VAE + 0.02 processor) plus

app.py

@@ -5,17 +5,18 @@
 import importlib.util
 import os
 
-# Less allocator fragmentation: the largest calls run within ~3.5 GiB of the card (see generate()).
+# Less allocator fragmentation (untiled 2048^2 text-to-image reserves ~65 GiB of the 96 GB card).
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 if importlib.util.find_spec("spaces"):
     import spaces
 
-    # size="large" (48 GB, 1x quota) is enough once the VAE decode is tiled above the 1024^2 bucket and for
-    # every editing call: the largest call (three-reference edit at 1760x1344) then peaks at 41.1 GiB
-    # allocated, versus 57.0 / 64.7 GiB for an untiled 2048^2 text-to-image (measured, B200, bf16 - see README).
-    # duration=90 covers the prompt rewrite (4-11 s on a B200, slower here) plus the largest call.
-    gpu = spaces.GPU(duration=90, size="large")
+    # size="xlarge" (the full 96 GB RTX Pro 6000 Blackwell, 2x quota): the pipeline runs exactly as validated, with no VAE
+    # tiling anywhere - tiling the encode wrecks reference-conditioned edits (duplicated subjects, wrong scale), and on the
+    # 48 GB "large" card (44.7 GiB usable) the untiled calls do not fit: three-reference edit at a 1024^2-area output
+    # 44.3 GiB, 2048^2 text-to-image 57.0 GiB, three-reference edit at 1760x1344 51.5 GiB (measured, B200, bf16 - see README).
+    # duration=90 covers the prompt rewrite plus the largest call.
+    gpu = spaces.GPU(duration=90, size="xlarge")
 else:
     gpu = lambda fn: fn  # noqa: E731
 
@@ -78,7 +79,6 @@ T2I_CHOICES = [*_bucket(1024), *_bucket(2048)]
 EDIT_CHOICES = [AUTO, *_bucket(1024), *_bucket(1536)]
 # the (width, height) pairs the editing menu actually offers, in menu order
 EDIT_DIMS = [SIZES[label] for label in EDIT_CHOICES if SIZES[label]]
-TILE_ABOVE_PIXELS = 1_500_000  # between the 1024² bucket (≤ 1_056_768 px) and the 1536² bucket (≥ 2_359_296 px)
 
 if STUDENT == "full":
     transformer = QwenImage21Transformer2DModel.from_pretrained(
@@ -95,6 +95,7 @@ else:
 pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift_terminal=None)
 pipe.to("cuda")
 
+
 # transformers' Qwen3VLVisionPatchEmbed runs an nn.Conv3d whose kernel_size equals its stride, which is
 # exactly a linear map over each flattened patch. cuDNN has no usable bf16 Conv3d kernel for this shape and
 # falls back to one that costs ~30 s per reference image (measured; 356 ms even with memory to spare, versus
@@ -177,17 +178,8 @@ def generate(prompt, image_1, image_2, image_3, size_label, seed, randomize_seed
         enhance_note = f" · enhance {time.perf_counter() - started:.1f}s"
         if used_prompt == prompt:
             enhance_note += " (rewrite failed to parse, original prompt used)"
-    # VAE tiling wherever the untiled decode would not fit the 48 GB (44.7 GiB) card: the 1536²/2048² buckets,
-    # and every editing call - with references resident the untiled decode of a 1024²-area output peaks at
-    # 44.3 GiB allocated for three references (42.0 for two), which is the OOM the first Space example hit
-    # (surfacing on ZeroGPU as "NVML_SUCCESS == r INTERNAL ASSERT FAILED" from the allocator). Tiled, the same
-    # call peaks at 40.3 GiB and the largest editing call (three references, 1760x1344) at 41.1 GiB, its
-    # denoising peak. Text-to-image at the 1024² bucket (38.2 GiB) stays untiled and bit-identical to the
-    # untiled pipeline; tiled decodes differ from untiled ones by 1.6/255 mean absolute (measured at 2048²).
-    if images or (width and width * height > TILE_ABOVE_PIXELS):
-        pipe.vae.enable_tiling()
-    else:
-        pipe.vae.disable_tiling()
+    # No VAE tiling anywhere (xlarge card): tiling the reference encode wrecks edits, and tiled decodes are not the
+    # validated pipeline either.
     generator = torch.Generator(device="cuda").manual_seed(seed)
     started = time.perf_counter()
     result = pipe(