Redirect to the authors' Space dipta007/OracleZoom

LICENSE-OracleZoom

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2026 Shubhashis Roy Dipta, Sourajit Saha, Shaswati Saha, Nobin Sarwar
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

README.md

@@ -3,64 +3,24 @@ title: OracleZoom Super Resolution
 emoji: 🔎
 colorFrom: purple
 colorTo: gray
-sdk: gradio
-sdk_version: 5.50.0
-python_version: "3.12.12"
-app_file: app.py
-startup_duration_timeout: 45m
+sdk: static
+app_file: index.html
+pinned: false
 license: mit
-short_description: Recursive 4x zoom to 256x with OracleZoom
+short_description: Moved — use the authors' Space dipta007/OracleZoom
 tags:
   - super-resolution
   - image-to-image
-  - diffusion
-  - zoom
 models:
   - dipta007/OracleZoom
-  - stabilityai/stable-diffusion-3-medium-diffusers
-  - Qwen/Qwen2.5-VL-3B-Instruct
-datasets:
-  - dipta007/OracleZoom-4KLSDB-train
 ---
 
-# OracleZoom 🔎
+# OracleZoom has moved 🔎
 
-Upload a photo, pick a point, and zoom into it 4x at a time. Each step super-resolves the
-*previous step's own output*, so four steps reach **256x**. A Qwen2.5-VL prompter reads the
-pair (current level, next blurry crop) and writes the conditioning for the SD3-medium
-one-step SR transformer that OracleZoom fine-tuned.
+This Space is now a redirect to the authors' own demo:
+**https://huggingface.co/spaces/dipta007/OracleZoom**
 
 - Paper: https://huggingface.co/papers/2609.06490
 - Code: https://github.com/dipta007/OracleZoom
 - Project page: https://dipta007.github.io/OracleZoom/
 - Model: https://huggingface.co/dipta007/OracleZoom
-- Training data: https://huggingface.co/datasets/dipta007/OracleZoom-4KLSDB-train
-- Authors' own Space: https://huggingface.co/spaces/dipta007/OracleZoom
-
-## What is in here
-
-`app.py` is the interface. The pipeline itself is the authors' inference path, reused
-unchanged from `space/` in [dipta007/OracleZoom](https://github.com/dipta007/OracleZoom)
-(MIT, see `LICENSE-OracleZoom`): `zoom.py` is the recursion, `geometry.py` the crop
-geometry, `video.py` the clip renderer, and `vendor/` holds the four files copied from
-[Chain-of-Zoom](https://github.com/bryanswkim/Chain-of-Zoom) (MIT, see `vendor/COZ-LICENSE`).
-
-The example images in `samples/` are the authors' own, shipped in the MIT-licensed repo and
-drawn from their [4KLSDB set](https://huggingface.co/datasets/dipta007/OracleZoom-4KLSDB-train)
-(CC-BY-4.0) — credit to Shubhashis Roy Dipta, Sourajit Saha, Shaswati Saha and Nobin Sarwar.
-
-## Notes for anyone duplicating this
-
-- Needs an `HF_TOKEN` secret: Stable Diffusion 3-medium is gated, so the token must come from
-  an account that accepted its licence. `huggingface_hub` picks the variable up on its own.
-- Hardware: ZeroGPU. The pipeline holds roughly 27 GB of weights, inside the 48 GB a `large`
-  slice gives. Boot streams SD3-medium, Qwen2.5-VL-3B and the merged transformer, so the
-  startup timeout is raised to 45 minutes.
-- Gradio is pinned to the 5.x series on purpose: `transformers==4.49.0` (what this pipeline
-  was built against) requires `huggingface-hub<1.0`, which gradio 6 cannot satisfy.
-
-## Honest limits
-
-Past the first step or two there is no ground truth left to recover, so the deep levels are
-plausible detail rather than measured detail. The zoom is a crop of your photo, not a real
-lens moving closer. The paper is explicit about which claims it does and does not make.

app.py

@@ -1,271 +0,0 @@
-"""OracleZoom — reference-constrained recursive image super-resolution.
-
-Upload a photo, pick a point, and zoom into it 4x at a time. Each step super-resolves the
-previous step's own output, so four steps reach 256x. A Qwen2.5-VL prompter reads the pair
-(current level, next blurry crop) and writes the conditioning for the SD3 one-step SR model.
-
-The pipeline is the authors' own inference path, taken from
-https://github.com/dipta007/OracleZoom (MIT): `zoom.py`, `geometry.py`, `video.py` and
-`vendor/` are theirs; only the interface below is new.
-"""
-import os
-import time
-
-# High-res pixel-space VAE decodes spike the allocator on ZeroGPU. Must precede any torch import.
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-# The VLM prompter tokenizes in the main process and every request forks a GPU worker.
-os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-
-import spaces  # noqa: E402  — before anything that touches torch/CUDA
-import gradio as gr  # noqa: E402
-from PIL import Image, ImageDraw  # noqa: E402
-
-import geometry  # noqa: E402
-import video  # noqa: E402
-import zoom  # noqa: E402
-
-UPSCALE = 4
-DEFAULT_STEPS = 4
-MAX_STEPS = 8
-ACCENT = "#a855f7"
-SAMPLES = ["0479", "0064", "0245", "0393", "0457"]
-
-# A real run, shipped with the Space, so the page shows the payoff before anyone spends quota.
-EXAMPLE_INPUT = "samples/example_1x.png"
-EXAMPLE_CLIP = "samples/example_zoom.mp4"
-EXAMPLE_LEVELS = [(EXAMPLE_INPUT, "input")] + \
-                 [(f"samples/example_{f}x.png", f"{f}x") for f in (4, 16, 64, 256)]
-EXAMPLE_COMPARE = ("samples/example_256x_input.png", "samples/example_256x.png")
-EXAMPLE_NOTE = "_From the shipped example run. Zoom your own photo to see its prompts._"
-
-MODELS = zoom.Models()
-
-HEADER = """# OracleZoom 🔎
-
-Zoom into a photo far past what it holds. Every step is a 4x super-resolution of the
-*previous step's own output*, so four steps reach **256x** — and the last available
-reference keeps each step honest about the structure it can still verify.
-
-[Paper](https://huggingface.co/papers/2609.06490) ·
-[Code](https://github.com/dipta007/OracleZoom) ·
-[Project page](https://dipta007.github.io/OracleZoom/) ·
-[Model](https://huggingface.co/dipta007/OracleZoom) ·
-[Training data](https://huggingface.co/datasets/dipta007/OracleZoom-4KLSDB-train)
-"""
-
-BIBTEX = """@misc{dipta2026oraclezoomonpolicyselfdistillationinspired,
-  title         = {OracleZoom: On-Policy Self-Distillation Inspired Reference-Constrained Recursive Image Super Resolution},
-  author        = {Shubhashis Roy Dipta and Sourajit Saha and Shaswati Saha and Nobin Sarwar},
-  year          = {2026},
-  eprint        = {2609.06490},
-  archivePrefix = {arXiv},
-  primaryClass  = {cs.CV},
-  url           = {https://arxiv.org/abs/2609.06490}
-}"""
-
-CSS = """
-#col-container {max-width: 1400px; margin: 0 auto;}
-.dark .gradio-container {color: var(--body-text-color);}
-#zoom-target img {cursor: crosshair;}
-#clip video {border-radius: 12px;}
-/* The level strip should hug its content: .gallery-container carries the height, so styling
-   only the outer block leaves a dead band below the thumbnails. */
-#levels, #levels .gallery-container, #levels .grid-wrap {
-  height: auto !important; max-height: none !important; min-height: 0 !important;
-  overflow: visible !important; flex-grow: 0 !important;}
-"""
-
-
-def label(factor: int) -> str:
-    return "input" if factor == 1 else f"{factor}x"
-
-
-def scale_labels(levels: int) -> list:
-    return [label(UPSCALE ** n) for n in range(1, int(levels) + 1)]
-
-
-def preview(image, levels, cx, cy):
-    """Draw the nested zoom windows so nobody spends a run to find out where it lands.
-
-    Falls back to the shipped example when nothing is uploaded yet — returning None would
-    blank the box the moment a visitor touched a slider.
-    """
-    canvas = geometry.resize_and_center_crop(
-        image if image is not None else Image.open(EXAMPLE_INPUT)).convert("RGB")
-    rects = geometry.nested_rects(canvas.size, int(levels), UPSCALE, (cx, cy))
-    dimmed = Image.blend(canvas, Image.new("RGB", canvas.size, (0, 0, 0)), 0.45)
-    if rects:
-        dimmed.paste(canvas.crop(rects[0]), rects[0][:2])
-    draw = ImageDraw.Draw(dimmed)
-    for rect in rects:
-        if rect[2] - rect[0] >= 3:
-            draw.rectangle(rect, outline=ACCENT, width=2)
-    return dimmed
-
-
-def pick_point(evt: gr.SelectData):
-    """Move the zoom point to wherever the preview was clicked (rounded to the slider step)."""
-    x, y = evt.index
-    return (round(x / geometry.PROCESS_SIZE, 2), round(y / geometry.PROCESS_SIZE, 2))
-
-
-def gallery_shape(levels):
-    """One row up to five levels, two beyond. Height is left to CSS, which hugs the content."""
-    items = int(levels) + 1
-    rows = 1 if items <= 5 else 2
-    return gr.update(columns=-(-items // rows), rows=rows)
-
-
-def _gallery_items(gallery):
-    """Gradio hands a filepath-typed Gallery back to a handler as (path, caption) tuples."""
-    return [(item[0], item[1] or "") for item in gallery or []
-            if isinstance(item, (tuple, list)) and len(item) == 2]
-
-
-def compare_level(gallery, cx, cy, evt: gr.SelectData):
-    """Rebuild the before/after for whichever level was clicked in the strip.
-
-    The left side is the plain bicubic enlargement that fed that step; recomputing it from the
-    levels already on screen keeps this on the CPU, so browsing costs no GPU quota.
-    """
-    items = _gallery_items(gallery)
-    index = getattr(evt, "index", None)
-    i = index if isinstance(index, int) else None
-    if i is None or i <= 0 or i >= len(items):
-        return gr.skip()
-    prev = Image.open(items[i - 1][0]).convert("RGB")
-    blurry = geometry.zoom_window(prev, UPSCALE, (cx, cy)).resize(prev.size, Image.BICUBIC)
-    return blurry, Image.open(items[i][0]).convert("RGB")
-
-
-def estimate_duration(image=None, levels=DEFAULT_STEPS, *args, **kwargs):
-    """GPU seconds to reserve. Checked against the visitor's remaining quota before the run,
-    so a loose number locks out low-quota visitors and a tight one gets the run killed.
-
-    Measured on this Space (warm worker, 512x512): 3.8 s at 1 step, 7.9 s at 4, 15.6 s at 8
-    — about 2 s fixed plus 1.8 s per step, clip rendering included. Doubled, which is what
-    covers a cold worker streaming the ~27 GB of weights into VRAM first.
-    """
-    return int(6 + 3 * int(levels or DEFAULT_STEPS))
-
-
-def _zoom_stream(image, levels, cx, cy):
-    if image is None:
-        raise gr.Error("Upload a photo first (or click one of the examples).")
-    levels = int(levels)
-    started = time.perf_counter()
-    gallery, prompts, frames, compare = [], [], [], None
-    for step, factor, prompt, blurry, result in zoom.zoom(MODELS, image, levels, UPSCALE,
-                                                          (cx, cy)):
-        frames.append(result)
-        gallery.append((result, label(factor)))
-        if step:
-            compare = (blurry, result)
-            prompts.append(f"**{label(factor)}** &nbsp; {prompt or '_(no prompt)_'}")
-        left = levels - step
-        if not step:
-            status = f"Warmed up. Zooming {levels} step{'s' if levels > 1 else ''}…"
-        elif left:
-            status = f"At **{label(factor)}**, {left} step{'s' if left > 1 else ''} to go…"
-        else:
-            status = f"At **{label(factor)}**. Rendering the clip…"
-        # The first yield clears any clip left from the previous run; later ones leave it alone.
-        yield (status, (None if not step else gr.skip()), gallery, compare,
-               "\n\n".join(prompts))
-    clip = video.render(frames, UPSCALE, (cx, cy))
-    elapsed = time.perf_counter() - started
-    print(f"#### zoom: {levels} step(s) in {elapsed:.1f}s", flush=True)
-    yield (f"Done. Reached **{label(UPSCALE ** levels)}** in {elapsed:.0f}s. "
-           f"Click any level below to compare it with the plain enlargement it started from.",
-           clip, gallery, compare, "\n\n".join(prompts))
-
-
-@spaces.GPU(duration=estimate_duration)
-def zoom_in(image, levels: int = DEFAULT_STEPS, cx: float = 0.5, cy: float = 0.5):
-    """Recursively super-resolve an image, 4x per step, streaming each level as it lands.
-
-    Args:
-        image: the photo to zoom into (resized and centre-cropped to 512x512 first).
-        levels: how many 4x steps to take. 4 steps reach 256x.
-        cx: horizontal position of the zoom point, 0 (left) to 1 (right).
-        cy: vertical position of the zoom point, 0 (top) to 1 (bottom).
-    """
-    yield from _zoom_stream(image, levels, cx, cy)
-
-
-@spaces.GPU(duration=estimate_duration)
-def zoom_in_cached(image, levels: int = DEFAULT_STEPS, cx: float = 0.5, cy: float = 0.5):
-    """Same zoom, returned in one shot so `gr.Examples` can cache it.
-
-    Args:
-        image: the photo to zoom into.
-        levels: how many 4x steps to take.
-        cx: horizontal position of the zoom point, 0 to 1.
-        cy: vertical position of the zoom point, 0 to 1.
-    """
-    final = None
-    for final in _zoom_stream(image, levels, cx, cy):
-        pass
-    return final
-
-
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="OracleZoom: zoom past 256x") as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(HEADER)
-        with gr.Row(equal_height=False):
-            with gr.Column(scale=4):
-                image = gr.Image(label="Your photo", type="pil", height=260,
-                                 sources=["upload", "clipboard"])
-                # Seeded with the shipped example so the overlay explains itself before any upload.
-                target = gr.Image(value=preview(None, DEFAULT_STEPS, 0.5, 0.5),
-                                  elem_id="zoom-target", label="Click to move the zoom point",
-                                  type="pil", interactive=False, height=330,
-                                  show_download_button=False)
-                levels = gr.Slider(1, MAX_STEPS, value=DEFAULT_STEPS, step=1, label="Zoom steps",
-                                   info=f"Each step is {UPSCALE}x. 4 steps reach 256x, "
-                                        f"{MAX_STEPS} reach {UPSCALE ** MAX_STEPS}x. "
-                                        f"Fewer steps finish sooner.")
-                go = gr.Button("🔎 Zoom in", variant="primary", size="lg")
-                status = gr.Markdown("Showing a shipped example run. "
-                                     "Upload a photo, then press **Zoom in**.")
-                with gr.Accordion("Advanced settings", open=False):
-                    cx = gr.Slider(0, 1, value=0.5, step=0.01, label="Zoom point — horizontal")
-                    cy = gr.Slider(0, 1, value=0.5, step=0.01, label="Zoom point — vertical")
-                gr.Markdown("##### What the VLM prompter read at each step")
-                prompts = gr.Markdown(EXAMPLE_NOTE)
-            with gr.Column(scale=6):
-                clip = gr.Video(value=EXAMPLE_CLIP, label="The zoom", elem_id="clip",
-                                autoplay=True, loop=True, height=460)
-                gallery = gr.Gallery(value=EXAMPLE_LEVELS, label="Every level, start to finish",
-                                     elem_id="levels", columns=5, rows=1, object_fit="cover",
-                                     allow_preview=False, show_download_button=True)
-                compare = gr.ImageSlider(value=EXAMPLE_COMPARE, height=470, elem_id="compare",
-                                         label="Plain enlargement (left) vs OracleZoom (right)")
-
-        gr.Examples(
-            examples=[f"samples/{n}.png" for n in SAMPLES],
-            inputs=[image],
-            outputs=[status, clip, gallery, compare, prompts],
-            fn=zoom_in_cached,
-            cache_examples=True,
-            cache_mode="lazy",
-            label="The authors' own example photos — click one to zoom it to 256x",
-        )
-
-        gr.Markdown("**Honest limits.** Past the first step or two there is no ground truth "
-                    "left to recover, so deep levels are plausible detail, not measured "
-                    "detail. The zoom is a crop of your photo, not a lens moving closer.")
-        gr.Markdown("### Cite")
-        gr.Code(value=BIBTEX, language=None, show_label=False, container=False)
-
-    controls = [image, levels, cx, cy]
-    for control in controls:
-        control.change(preview, controls, target, show_api=False)
-    target.select(pick_point, None, [cx, cy], show_api=False)
-    # Reshape before the run, not after, so the strip does not resize under the results.
-    levels.change(gallery_shape, levels, gallery, show_api=False)
-    go.click(zoom_in, controls, [status, clip, gallery, compare, prompts], api_name="zoom_in")
-    gallery.select(compare_level, [gallery, cx, cy], compare, show_api=False)
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(mcp_server=True)

geometry.py

@@ -1,45 +0,0 @@
-"""Crop geometry for the zoom. Pure PIL, no model code, so it stays cheap to test."""
-from PIL import Image
-
-PROCESS_SIZE = 512
-
-
-def resize_and_center_crop(img, size=PROCESS_SIZE):
-    """CoZ's entry crop: shortest side to `size`, then the centre square."""
-    w, h = img.size
-    scale = size / min(w, h)
-    nw, nh = int(w * scale), int(h * scale)
-    img = img.resize((nw, nh), Image.LANCZOS)
-    left, top = (nw - size) // 2, (nh - size) // 2
-    return img.crop((left, top, left + size, top + size))
-
-
-def window_rect(size, upscale, center=(0.5, 0.5)):
-    """The next level's field of view: a 1/upscale square around `center`, clamped inside."""
-    w, h = size
-    nw, nh = w // upscale, h // upscale
-    left = min(max(int(center[0] * w - nw / 2), 0), w - nw)
-    top = min(max(int(center[1] * h - nh / 2), 0), h - nh)
-    return left, top, left + nw, top + nh
-
-
-def zoom_window(img, upscale, center=(0.5, 0.5)):
-    return img.crop(window_rect(img.size, upscale, center))
-
-
-def nested_rects(size, levels, upscale=4, center=(0.5, 0.5)):
-    """Where each zoom step lands, in the coordinates of the first canvas.
-
-    Stops early once a window would be under a pixel wide. The recursion itself never
-    degenerates, because every step resizes its crop back up to the full canvas, but drawing
-    all those steps onto the one original canvas does run out of pixels.
-    """
-    rects, rect = [], (0, 0, size[0], size[1])
-    for _ in range(levels):
-        side = rect[2] - rect[0]
-        if side // upscale < 1:
-            break
-        l, t, r, b = window_rect((side, side), upscale, center)
-        rect = (rect[0] + l, rect[1] + t, rect[0] + r, rect[1] + b)
-        rects.append(rect)
-    return rects

index.html

@@ -0,0 +1,93 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>OracleZoom has moved</title>
+<link rel="canonical" href="https://huggingface.co/spaces/dipta007/OracleZoom">
+<style>
+  :root {
+    color-scheme: light dark;
+    --bg: #ffffff;
+    --fg: #111827;
+    --muted: #6b7280;
+    --line: #e5e7eb;
+    --accent: #ff9d00;
+  }
+  @media (prefers-color-scheme: dark) {
+    :root { --bg: #0b0f19; --fg: #f3f4f6; --muted: #9ca3af; --line: #1f2937; }
+  }
+  * { box-sizing: border-box; }
+  body {
+    margin: 0;
+    min-height: 100vh;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    padding: 24px;
+    background: var(--bg);
+    color: var(--fg);
+    font-family: ui-sans-serif, system-ui, -apple-system, "Segoe UI", Roboto, sans-serif;
+    line-height: 1.55;
+  }
+  main { max-width: 34rem; text-align: center; }
+  .emoji { font-size: 3rem; line-height: 1; }
+  h1 { font-size: 1.5rem; font-weight: 650; margin: 1rem 0 0.5rem; }
+  p { color: var(--muted); margin: 0 0 1.5rem; }
+  a.cta {
+    display: inline-block;
+    padding: 0.75rem 1.5rem;
+    border-radius: 0.65rem;
+    background: var(--accent);
+    color: #111827;
+    font-weight: 600;
+    text-decoration: none;
+  }
+  a.cta:hover { filter: brightness(1.05); }
+  .links {
+    margin-top: 2rem;
+    padding-top: 1.25rem;
+    border-top: 1px solid var(--line);
+    font-size: 0.875rem;
+    color: var(--muted);
+  }
+  .links a { color: inherit; }
+  .count { font-variant-numeric: tabular-nums; }
+</style>
+</head>
+<body>
+<main>
+  <div class="emoji">🔎</div>
+  <h1>OracleZoom lives on the authors' Space</h1>
+  <p>
+    Redirecting to <strong>dipta007/OracleZoom</strong> in <span class="count" id="count">3</span>s.
+  </p>
+  <a class="cta" id="go" href="https://huggingface.co/spaces/dipta007/OracleZoom" target="_top" rel="noopener">
+    Open OracleZoom →
+  </a>
+  <div class="links">
+    <a href="https://huggingface.co/papers/2609.06490" target="_blank" rel="noopener">Paper</a> ·
+    <a href="https://github.com/dipta007/OracleZoom" target="_blank" rel="noopener">Code</a> ·
+    <a href="https://dipta007.github.io/OracleZoom/" target="_blank" rel="noopener">Project page</a> ·
+    <a href="https://huggingface.co/dipta007/OracleZoom" target="_blank" rel="noopener">Model</a>
+  </div>
+</main>
+<script>
+  const TARGET = "https://huggingface.co/spaces/dipta007/OracleZoom";
+  const el = document.getElementById("count");
+  let n = 3;
+  const tick = setInterval(() => {
+    n -= 1;
+    if (n > 0) { el.textContent = n; return; }
+    clearInterval(tick);
+    el.textContent = 0;
+    try {
+      if (window.top !== window.self) { window.top.location.replace(TARGET); }
+      else { window.location.replace(TARGET); }
+    } catch (e) {
+      document.getElementById("go").focus();
+    }
+  }, 1000);
+</script>
+</body>
+</html>

requirements.txt

@@ -1,28 +0,0 @@
-# Pins taken from the authors' own Space (github.com/dipta007/OracleZoom, space/requirements.txt).
-# diffusers stays at the paper's version: the SR LoRA is injected by matching the class name
-# "AdaLayerNormZero", so a diffusers rename would silently inject nothing.
-# transformers 4.49.0 is what the Qwen2.5-VL prompter + peft merge path was built against; it
-# also pins huggingface-hub<1.0, which is why the Space runs gradio 5.50 rather than 6.x.
-# gradio / spaces / huggingface_hub are deliberately absent (platform-managed).
-torch==2.8.0
-torchvision==0.23.0
-diffusers==0.32.2
-transformers==4.49.0
-accelerate
-peft>=0.13
-safetensors
-qwen-vl-utils
-# SD3's third text encoder is T5. Its fast tokenizer is built from a sentencepiece model, and
-# the conversion wants protobuf. Both arrive transitively in the research env, not here.
-sentencepiece
-protobuf
-lpips
-einops
-pyyaml
-tqdm
-numpy
-pillow
-imageio
-imageio-ffmpeg
-# demo.launch(mcp_server=True)
-mcp

samples/example_zoom.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dc9b6172522eb830313cf1f18f9bb852706886005850dffd80e45a66f91391c7
-size 3691325

vendor/COZ-LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2025 Bryan Sangwoo Kim
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

vendor/lora/lora_layers.py

@@ -1,137 +0,0 @@
-from typing import List, Optional, Set, Type, Union
-
-import torch
-from torch import nn
-
-
-class LoraInjectedLinear(nn.Module):
-    """
-    Linear layer with LoRA injection.
-    Taken from https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
-    """
-    def __init__(
-        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
-    ):
-        super().__init__()
-
-        if r > min(in_features, out_features):
-            raise ValueError(
-                f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
-            )
-        self.r = r
-        self.linear = nn.Linear(in_features, out_features, bias)
-        self.lora_down = nn.Linear(in_features, r, bias=False)
-        self.dropout = nn.Dropout(dropout_p)
-        self.lora_up = nn.Linear(r, out_features, bias=False)
-        self.scale = scale
-        self.selector = nn.Identity()
-
-        nn.init.normal_(self.lora_down.weight, std=1 / r)
-        nn.init.zeros_(self.lora_up.weight)
-
-    def forward(self, input):
-        return (
-            self.linear(input.float())
-            + self.dropout(self.lora_up(self.selector(self.lora_down(input.float()))))
-            * self.scale
-        ).half()
-
-    def realize_as_lora(self):
-        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
-
-    def set_selector_from_diag(self, diag: torch.Tensor):
-        # diag is a 1D tensor of size (r,)
-        assert diag.shape == (self.r,)
-        self.selector = nn.Linear(self.r, self.r, bias=False)
-        self.selector.weight.data = torch.diag(diag)
-        self.selector.weight.data = self.selector.weight.data.to(
-            self.lora_up.weight.device
-        ).to(self.lora_up.weight.dtype)
-
-class LoraInjectedConv2d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups: int = 1,
-        bias: bool = True,
-        r: int = 4,
-        dropout_p: float = 0.1,
-        scale: float = 1.0,
-    ):
-        super().__init__()
-        if r > min(in_channels, out_channels):
-            raise ValueError(
-                f"LoRA rank {r} must be less or equal than {min(in_channels, out_channels)}"
-            )
-        self.r = r
-        self.conv = nn.Conv2d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=bias,
-        )
-
-        self.lora_down = nn.Conv2d(
-            in_channels=in_channels,
-            out_channels=r,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=False,
-        )
-        self.dropout = nn.Dropout(dropout_p)
-        self.lora_up = nn.Conv2d(
-            in_channels=r,
-            out_channels=out_channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.selector = nn.Identity()
-        self.scale = scale
-
-        nn.init.normal_(self.lora_down.weight, std=1 / r)
-        nn.init.zeros_(self.lora_up.weight)
-
-    def forward(self, input):
-        return (
-            self.conv(input)
-            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
-            * self.scale
-        )
-
-    def realize_as_lora(self):
-        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
-
-    def set_selector_from_diag(self, diag: torch.Tensor):
-        # diag is a 1D tensor of size (r,)
-        assert diag.shape == (self.r,)
-        self.selector = nn.Conv2d(
-            in_channels=self.r,
-            out_channels=self.r,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.selector.weight.data = torch.diag(diag)
-
-        # same device + dtype as lora_up
-        self.selector.weight.data = self.selector.weight.data.to(
-            self.lora_up.weight.device
-        ).to(self.lora_up.weight.dtype)
-
-
-

vendor/lora/lora_utils.py

@@ -1,61 +0,0 @@
-import torch
-from torch import nn
-from lora.lora_layers import LoraInjectedLinear, LoraInjectedConv2d
-
-def _find_modules(model, ancestor_class=None, search_class=[nn.Linear], exclude_children_of=[LoraInjectedLinear]):
-    # Get the targets we should replace all linears under
-    if ancestor_class is not None:
-        ancestors = (
-            module
-            for module in model.modules()
-            if module.__class__.__name__ in ancestor_class
-        )
-    else:
-        # this, incase you want to naively iterate over all modules.
-        ancestors = [module for module in model.modules()]
-
-    for ancestor in ancestors:
-        for fullname, module in ancestor.named_modules():
-            # if 'norm1_context' in fullname:
-                if any([isinstance(module, _class) for _class in search_class]):
-                    *path, name = fullname.split(".")
-                    parent = ancestor
-                    while path:
-                        parent = parent.get_submodule(path.pop(0))
-                    if exclude_children_of and any(
-                        [isinstance(parent, _class) for _class in exclude_children_of]
-                    ):
-                        continue
-                    yield parent, name, module
-
-def extract_lora_ups_down(model, target_replace_module={'AdaLayerNormZero'}):   # Attention for kv_lora
-
-    loras = []
-
-    for _m, _n, _child_module in _find_modules(
-        model,
-        target_replace_module,
-        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
-    ):
-        loras.append((_child_module.lora_up, _child_module.lora_down))
-
-    if len(loras) == 0:
-        raise ValueError("No lora injected.")
-
-    return loras
-
-def save_lora_weight(
-    model,
-    path="./lora.pt",
-    target_replace_module={'AdaLayerNormZero'},     # Attention for kv_lora
-    save_half:bool=False
-):
-    weights = []
-    for _up, _down in extract_lora_ups_down(
-        model, target_replace_module=target_replace_module
-    ):
-        dtype = torch.float16 if save_half else torch.float32
-        weights.append(_up.weight.to("cpu").to(dtype))
-        weights.append(_down.weight.to("cpu").to(dtype))
-
-    torch.save(weights, path)

vendor/osediff_sd3.py

@@ -1,913 +0,0 @@
-import os
-import sys
-sys.path.append(os.getcwd())
-import yaml
-import copy
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import List, Tuple, Optional
-import numpy as np
-import lpips
-from torchvision import transforms
-from PIL import Image
-from peft import LoraConfig, get_peft_model
-
-from copy import deepcopy
-from tqdm import tqdm
-
-from diffusers import StableDiffusion3Pipeline
-from lora.lora_layers import LoraInjectedLinear, LoraInjectedConv2d
-
-from utils.vaehook import VAEHook
-
-
-def inject_lora_vae(vae, lora_rank=4, init_lora_weights="gaussian", verbose=False):
-    """
-    Inject LoRA into the VAE's encoder
-    """
-    vae.requires_grad_(False)
-    vae.train() 
-
-    # Identify modules to LoRA-ify in the encoder
-    l_grep = ["conv1", "conv2", "conv_in", "conv_shortcut", 
-              "conv", "conv_out", "to_k", "to_q", "to_v", "to_out.0"]
-    l_target_modules_encoder = []
-    for n, p in vae.named_parameters():
-        if "bias" in n or "norm" in n:
-            continue
-        for pattern in l_grep:
-            if (pattern in n) and ("encoder" in n):
-                l_target_modules_encoder.append(n.replace(".weight", ""))
-            elif ("quant_conv" in n) and ("post_quant_conv" not in n):
-                l_target_modules_encoder.append(n.replace(".weight", ""))
-
-    if verbose:
-        print("The following VAE parameters will get LoRA:")
-        print(l_target_modules_encoder)
-
-    # Create and add a LoRA adapter
-    lora_conf_encoder = LoraConfig(
-        r=lora_rank,
-        init_lora_weights=init_lora_weights,
-        target_modules=l_target_modules_encoder
-    )
-    
-    adapter_name = "default_encoder"
-    try:
-        vae.add_adapter(lora_conf_encoder, adapter_name=adapter_name)
-        vae.set_adapter(adapter_name)
-    except ValueError as e:
-        if "already exists" in str(e):
-            print(f"Adapter with name {adapter_name} already exists. Skipping injection.")
-        else:
-            raise e
-
-    return vae, l_target_modules_encoder
-
-def _find_modules(model, ancestor_class=None, search_class=[nn.Linear], exclude_children_of=[LoraInjectedLinear]):
-    # Get the targets we should replace all linears under
-    if ancestor_class is not None:
-        ancestors = (
-            module
-            for module in model.modules()
-            if module.__class__.__name__ in ancestor_class
-        )
-    else:
-        # this, in case you want to naively iterate over all modules.
-        ancestors = [module for module in model.modules()]
-
-    for ancestor in ancestors:
-        for fullname, module in ancestor.named_modules():
-            if any([isinstance(module, _class) for _class in search_class]):
-                *path, name = fullname.split(".")
-                parent = ancestor
-                while path:
-                    parent = parent.get_submodule(path.pop(0))
-                if exclude_children_of and any(
-                    [isinstance(parent, _class) for _class in exclude_children_of]
-                ):
-                    continue
-                yield parent, name, module
-
-def inject_lora(model, ancestor_class, loras=None, r:int=4, dropout_p:float=0.0, scale:float=1.0, verbose:bool=False):
-    
-    model.requires_grad_(False)
-    model.train()
-    
-    names = []
-    require_grad_params = []  # to be updated
-
-    total_lora_params = 0  
-
-    if loras is not None:
-        loras = torch.load(loras, map_location=model.device, weights_only=True)
-        loras = [lora.float() for lora in loras]
-
-    for _module, name, _child_module in _find_modules(model, ancestor_class):  # SiLU + Linear Block
-        weight = _child_module.weight
-        bias = _child_module.bias
-
-        if verbose:
-            print(f'LoRA Injection : injecting lora into {name}')
-
-        _tmp = LoraInjectedLinear(
-            _child_module.in_features,
-            _child_module.out_features,
-            _child_module.bias is not None,
-            r=r,
-            dropout_p=dropout_p,
-            scale=scale,
-        )
-        _tmp.linear.weight = nn.Parameter(weight.float())
-        if bias is not None:
-            _tmp.linear.bias = nn.Parameter(bias.float())
-        
-        # switch the module
-        _tmp.to(device=_child_module.weight.device, dtype=torch.float) # keep as float / mixed precision
-        _module._modules[name] = _tmp
-
-        require_grad_params.append(_module._modules[name].lora_up.parameters())
-        require_grad_params.append(_module._modules[name].lora_down.parameters())
-
-        if loras != None:
-            _module._modules[name].lora_up.weight = nn.Parameter(loras.pop(0))
-            _module._modules[name].lora_down.weight = nn.Parameter(loras.pop(0))
-
-        _module._modules[name].lora_up.weight.requires_grad = True
-        _module._modules[name].lora_down.weight.requires_grad = True
-        names.append(name)
-
-        if verbose:
-            # -------- Count LoRA parameters just added --------
-            lora_up_count = sum(p.numel() for p in _tmp.lora_up.parameters())
-            lora_down_count = sum(p.numel() for p in _tmp.lora_down.parameters())
-            lora_total_for_this_layer = lora_up_count + lora_down_count
-            total_lora_params += lora_total_for_this_layer
-            print(f"  Added {lora_total_for_this_layer} params "
-                  f"(lora_up={lora_up_count}, lora_down={lora_down_count})")
-    
-    if verbose:
-        print(f"Total new LoRA parameters added: {total_lora_params}")
-    
-    return require_grad_params, names
-
-def add_mp_hook(transformer):
-    '''
-    For mixed precision of LoRA. (i.e. keep LoRA as float and others as half)
-    '''
-    def pre_hook(module, input):
-        return input.float()
-    
-    def post_hook(module, input, output):
-        return output.half()
-
-    hooks = []
-    for _module, name, _child_module in _find_modules(transformer):
-        if isinstance(_child_module, LoraInjectedLinear):
-            hook = _child_module.lora_up.register_forward_pre_hook(pre_hook)
-            hooks.append(hook)
-            hook = _child_module.lora_down.register_forward_hook(post_hook)
-            hooks.append(hook)
-    
-    return transformer, hooks
-
-def compute_density_for_timestep_sampling(
-    weighting_scheme: str, batch_size: int, logit_mean: float = 0.0, logit_std: float = 1.0, mode_scale: Optional[float] = None
-):
-    """
-    Compute the density for sampling the timesteps when doing SD3 training.
-
-    Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
-
-    SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
-    """
-    if weighting_scheme == "logit_normal":
-        # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
-        u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu")
-        u = torch.nn.functional.sigmoid(u)
-    elif weighting_scheme == "mode":
-        u = torch.rand(size=(batch_size,), device="cpu")
-        u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
-    else:
-        u = torch.rand(size=(batch_size,), device="cpu")
-    return u
-
-def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas):
-    """
-    Computes loss weighting scheme for SD3 training.
-
-    Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
-
-    SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
-    """
-    if weighting_scheme == "sigma_sqrt":
-        weighting = (sigmas**-2.0).float()
-    elif weighting_scheme == "cosmap":
-        bot = 1 - 2 * sigmas + 2 * sigmas**2
-        weighting = 2 / (math.pi * bot)
-    else:
-        weighting = torch.ones_like(sigmas)
-    return weighting
-
-
-class StableDiffusion3Base():
-    def __init__(self, model_key:str='stabilityai/stable-diffusion-3-medium-diffusers', device='cuda', dtype=torch.float16):
-        self.device = device
-        self.dtype = dtype
-
-        pipe = StableDiffusion3Pipeline.from_pretrained(model_key, torch_dtype=self.dtype)
-
-        self.scheduler = pipe.scheduler
-
-        self.tokenizer_1 = pipe.tokenizer
-        self.tokenizer_2 = pipe.tokenizer_2
-        self.tokenizer_3 = pipe.tokenizer_3
-        self.text_enc_1 = pipe.text_encoder.to(device)
-        self.text_enc_2 = pipe.text_encoder_2.to(device)
-        self.text_enc_3 = pipe.text_encoder_3.to(device)
-
-        self.vae=pipe.vae.to(device)
-
-        self.transformer = pipe.transformer.to(device)
-        self.transformer.eval()
-        self.transformer.requires_grad_(False)
-
-        self.vae_scale_factor = (
-            2 ** (len(self.vae.config.block_out_channels)-1) if hasattr(self, "vae") and self.vae is not None else 8
-        )
-
-        del pipe
-
-    def encode_prompt(self, prompt: List[str], batch_size:int=1) -> List[torch.Tensor]:
-        '''
-        We assume that
-        1. number of tokens < max_length
-        2. one prompt for one image
-        '''
-        # CLIP encode (used for modulation of adaLN-zero)
-        # now, we have two CLIPs
-        text_clip1_ids = self.tokenizer_1(prompt,
-                                          padding="max_length",
-                                          max_length=77,
-                                          truncation=True,
-                                          return_tensors='pt').input_ids
-        text_clip1_emb = self.text_enc_1(text_clip1_ids.to(self.device), output_hidden_states=True)
-        pool_clip1_emb = text_clip1_emb[0].to(dtype=self.dtype, device=self.device)
-        text_clip1_emb = text_clip1_emb.hidden_states[-2].to(dtype=self.dtype, device=self.device)
-
-        text_clip2_ids = self.tokenizer_2(prompt,
-                                          padding="max_length",
-                                          max_length=77,
-                                          truncation=True,
-                                          return_tensors='pt').input_ids
-        text_clip2_emb = self.text_enc_2(text_clip2_ids.to(self.device), output_hidden_states=True)
-        pool_clip2_emb = text_clip2_emb[0].to(dtype=self.dtype, device=self.device)
-        text_clip2_emb = text_clip2_emb.hidden_states[-2].to(dtype=self.dtype, device=self.device)
-        
-        # T5 encode (used for text condition)
-        text_t5_ids = self.tokenizer_3(prompt,
-                                       padding="max_length",
-                                       max_length=512,
-                                       truncation=True,
-                                       add_special_tokens=True,
-                                       return_tensors='pt').input_ids
-        text_t5_emb = self.text_enc_3(text_t5_ids.to(self.device))[0]
-        text_t5_emb = text_t5_emb.to(dtype=self.dtype, device=self.device)
-
-        # Merge
-        clip_prompt_emb = torch.cat([text_clip1_emb, text_clip2_emb], dim=-1)
-        clip_prompt_emb = torch.nn.functional.pad(
-            clip_prompt_emb, (0, text_t5_emb.shape[-1] - clip_prompt_emb.shape[-1])
-        )
-        prompt_emb = torch.cat([clip_prompt_emb, text_t5_emb], dim=-2)
-        pooled_prompt_emb = torch.cat([pool_clip1_emb, pool_clip2_emb], dim=-1)
-
-        return prompt_emb, pooled_prompt_emb
-
-    def initialize_latent(self, img_size:Tuple[int], batch_size:int=1, **kwargs):
-        H, W = img_size
-        lH, lW = H//self.vae_scale_factor, W//self.vae_scale_factor
-        lC = self.transformer.config.in_channels
-        latent_shape = (batch_size, lC, lH, lW)
-
-        z = torch.randn(latent_shape, device=self.device, dtype=self.dtype)
-
-        return z
-
-    def encode(self, image: torch.Tensor) -> torch.Tensor:
-        z = self.vae.encode(image).latent_dist.sample()
-        z = (z-self.vae.config.shift_factor) * self.vae.config.scaling_factor
-        return z
-    
-    def decode(self, z: torch.Tensor) -> torch.Tensor:
-        z = (z/self.vae.config.scaling_factor) + self.vae.config.shift_factor
-        return self.vae.decode(z, return_dict=False)[0]
-    
-    
-class SD3Euler(StableDiffusion3Base):
-    def __init__(self, model_key:str='stabilityai/stable-diffusion-3-medium-diffusers', device='cuda'):
-        super().__init__(model_key=model_key, device=device)
-
-    def inversion(self, src_img, prompts: List[str], NFE:int, cfg_scale: float=1.0, batch_size: int=1):
-
-        # encode text prompts
-        prompt_emb, pooled_emb = self.encode_prompt(prompts, batch_size)
-        null_prompt_emb, null_pooled_emb = self.encode_prompt([""], batch_size)
-
-        # initialize latent
-        src_img = src_img.to(device=self.device, dtype=self.dtype)
-        with torch.no_grad():
-            z = self.encode(src_img)
-            z0 = z.clone()
-
-        # timesteps (default option. You can make your custom here.)
-        self.scheduler.set_timesteps(NFE, device=self.device)
-        timesteps = self.scheduler.timesteps
-        timesteps = torch.cat([timesteps, torch.zeros(1, device=self.device)])
-        timesteps = reversed(timesteps)
-        sigmas = timesteps / self.scheduler.config.num_train_timesteps
-
-        # Solve ODE
-        pbar = tqdm(timesteps[:-1], total=NFE, desc='SD3 Euler Inversion')
-        for i, t in enumerate(pbar):
-            timestep = t.expand(z.shape[0]).to(self.device)
-            pred_v = self.predict_vector(z, timestep, prompt_emb, pooled_emb)
-            if cfg_scale != 1.0:
-                pred_null_v = self.predict_vector(z, timestep, null_prompt_emb, null_pooled_emb)
-            else:
-                pred_null_v = 0.0
-
-            sigma = sigmas[i]
-            sigma_next = sigmas[i+1]
-
-            z = z + (sigma_next - sigma) * (pred_null_v + cfg_scale * (pred_v - pred_null_v))
-
-        return z
-
-    def sample(self, prompts: List[str], NFE:int, img_shape: Optional[Tuple[int]]=None, cfg_scale: float=1.0, batch_size: int = 1, latent:Optional[torch.Tensor]=None):
-        imgH, imgW = img_shape if img_shape is not None else (512, 512)
-
-        # encode text prompts
-        with torch.no_grad():
-            prompt_emb, pooled_emb = self.encode_prompt(prompts, batch_size)
-            null_prompt_emb, null_pooled_emb = self.encode_prompt([""], batch_size)
-
-        # initialize latent
-        if latent is None:
-            z = self.initialize_latent((imgH, imgW), batch_size)
-        else:
-            z = latent
-
-        # timesteps (default option. You can make your custom here.)
-        self.scheduler.set_timesteps(NFE, device=self.device)
-        timesteps = self.scheduler.timesteps
-        sigmas = timesteps / self.scheduler.config.num_train_timesteps
-
-        # Solve ODE
-        pbar = tqdm(timesteps, total=NFE, desc='SD3 Euler')
-        for i, t in enumerate(pbar):
-            timestep = t.expand(z.shape[0]).to(self.device)
-            pred_v = self.predict_vector(z, timestep, prompt_emb, pooled_emb)
-            if cfg_scale != 1.0:
-                pred_null_v = self.predict_vector(z, timestep, null_prompt_emb, null_pooled_emb)
-            else:
-                pred_null_v = 0.0
-
-            sigma = sigmas[i]
-            sigma_next = sigmas[i+1] if i+1 < NFE else 0.0
-
-            z = z + (sigma_next - sigma) * (pred_null_v + cfg_scale * (pred_v - pred_null_v))
-
-        # decode
-        with torch.no_grad():
-            img = self.decode(z)
-        return img
-
-
-class OSEDiff_SD3_GEN(torch.nn.Module):
-    def __init__(self, args, base_model):
-        super().__init__()
-        
-        self.args = args
-        self.model = base_model
-        
-        # Add lora to transformer
-        print('Adding LoRA to OSEDiff_SD3_GEN')
-        self.transformer_gen = copy.deepcopy(self.model.transformer)
-        self.transformer_gen.to('cuda:1')
-        
-        self.transformer_gen.requires_grad_(False)
-        self.transformer_gen.train()
-        self.transformer_gen, hooks = add_mp_hook(self.transformer_gen)
-        self.hooks = hooks
-
-        lora_params, _ = inject_lora(self.transformer_gen, {"AdaLayerNormZero"}, r=args.lora_rank, verbose=False)
-        for name, param in self.transformer_gen.named_parameters():
-            if "lora_" in name:
-                param.requires_grad = True   # LoRA up/down
-            else:
-                param.requires_grad = False  # everything else
-        
-        # Insert LoRA into VAE
-        print("Adding LoRA to VAE")
-        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
-
-    def predict_vector(self, z, t, prompt_emb, pooled_emb):
-        v = self.transformer_gen(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-
-    def forward(self, x_src, batch=None, args=None):
-        
-        z_src = self.model.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device))
-        z_src = z_src.to(self.transformer_gen.device)
-        
-        # calculate prompt_embeddings and neg_prompt_embeddings
-        batch_size, _, _, _ = x_src.shape
-        with torch.no_grad():
-            prompt_embeds, pooled_embeds = self.model.encode_prompt(batch["prompt"], batch_size)
-            neg_prompt_embeds, neg_pooled_embeds = self.model.encode_prompt(batch["neg_prompt"], batch_size)
-        
-        NFE = 1
-        self.model.scheduler.set_timesteps(NFE, device=self.model.device)
-        timesteps = self.model.scheduler.timesteps
-        sigmas = timesteps / self.model.scheduler.config.num_train_timesteps
-        sigmas = sigmas.to(self.transformer_gen.device)
-
-        # Solve ODE
-        i = 0
-        t = timesteps[0]
-
-        timestep = t.expand(z_src.shape[0]).to(self.transformer_gen.device)
-        prompt_embeds = prompt_embeds.to(self.transformer_gen.device, dtype=torch.float32)
-        pooled_embeds = pooled_embeds.to(self.transformer_gen.device, dtype=torch.float32)
-        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
-        pred_null_v = 0.0
-
-        sigma = sigmas[i]
-        sigma_next = sigmas[i+1] if i+1 < NFE else 0.0
-        
-        z_src = z_src + (sigma_next - sigma) * (pred_null_v + 1 * (pred_v - pred_null_v))
-        
-        output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
-
-        return output_image, z_src, prompt_embeds, pooled_embeds
-
-
-class OSEDiff_SD3_REG(torch.nn.Module):
-    def __init__(self, args, base_model):
-        super().__init__()
-        
-        self.args = args
-        self.model = base_model
-        self.transformer_org = self.model.transformer
-        
-        # Add lora to transformer
-        print('Adding LoRA to OSEDiff_SD3_REG')
-        self.transformer_reg = copy.deepcopy(self.transformer_org)
-        self.transformer_reg.to('cuda:1')
-        
-        self.transformer_reg.requires_grad_(False)
-        self.transformer_reg.train()
-        self.transformer_reg, hooks = add_mp_hook(self.transformer_reg)
-        self.hooks = hooks
-        
-        lora_params, _ = inject_lora(self.transformer_reg, {"AdaLayerNormZero"}, r=args.lora_rank, verbose=False)
-        for name, param in self.transformer_reg.named_parameters():
-            if "lora_" in name:
-                param.requires_grad = True   # LoRA up/down
-            else:
-                param.requires_grad = False  # everything else
-        
-    def predict_vector_reg(self, z, t, prompt_emb, pooled_emb):
-        v = self.transformer_reg(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-    
-    def predict_vector_org(self, z, t, prompt_emb, pooled_emb):
-        v = self.transformer_org(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-
-    def distribution_matching_loss(self, z0, prompt_embeds, pooled_embeds, global_step, args):
-        
-        with torch.no_grad():
-            device = self.transformer_reg.device
-            # get timesteps and sigma
-            u = compute_density_for_timestep_sampling(
-                        weighting_scheme="uniform",
-                        batch_size=1,
-                        logit_mean=0.0,
-                        logit_std=1.0,
-                        mode_scale=1.29,
-                    )
-            
-            t_idx = (u*1000).long().to(device)
-            self.model.scheduler.set_timesteps(1000, device=device)
-            times = self.model.scheduler.timesteps
-            t = times[t_idx]
-            sigma = t / 1000
-            
-            # get noise and xt
-            z0 = z0.to(device)
-            noise = torch.randn_like(z0)
-            sigma = sigma.half()
-            zt = (1-sigma) * z0 + sigma * noise
-            
-            # Get x0_prediction of transformer_reg
-            v_pred_reg = self.predict_vector_reg(zt, t, prompt_embeds.to(device), pooled_embeds.to(device))
-            reg_model_pred = v_pred_reg * (-sigma) + zt         # this is x0_prediction for reg
-            
-            # Get x0_prediction of transformer_org
-            org_device = self.transformer_org.device
-            v_pred_org = self.predict_vector_org(zt.to(org_device), t.to(org_device), prompt_embeds.to(org_device), pooled_embeds.to(org_device))
-            org_model_pred = v_pred_org * (-sigma.to(org_device)) + zt.to(org_device)         # this is x0_prediction for org
-            
-            # Visualization
-            if global_step % 100 == 1:
-                self.vsd_visualization(z0, noise, zt, reg_model_pred, org_model_pred, global_step, args)
-            
-        weighting_factor = torch.abs(z0 - org_model_pred.to(device)).mean(dim=[1, 2, 3], keepdim=True)
-
-        grad = (reg_model_pred - org_model_pred.to(device)) / weighting_factor
-        loss = F.mse_loss(z0, (z0 - grad).detach())
-
-        return loss
-    
-    def vsd_visualization(self, z0, noise, zt, reg_model_pred, org_model_pred, global_step, args):
-        #-------- Visualization --------#
-        # 1. Visualize latents, noise, zt
-        z0_img = self.model.decode(z0.to(dtype=torch.float32, device=self.model.vae.device))
-        ns_img = self.model.decode(noise.to(dtype=torch.float32, device=self.model.vae.device))
-        zt_img = self.model.decode(zt.to(dtype=torch.float32, device=self.model.vae.device))
-        
-        z0_img_pil = transforms.ToPILImage()(torch.clamp(z0_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
-        ns_img_pil = transforms.ToPILImage()(torch.clamp(ns_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
-        zt_img_pil = transforms.ToPILImage()(torch.clamp(zt_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
-        
-        # 2. Visualize reg_img, org_img
-        reg_img = self.model.decode(reg_model_pred.to(dtype=torch.float32, device=self.model.vae.device))
-        org_img = self.model.decode(org_model_pred.to(dtype=torch.float32, device=self.model.vae.device))
-        
-        reg_img_pil = transforms.ToPILImage()(torch.clamp(reg_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
-        org_img_pil = transforms.ToPILImage()(torch.clamp(org_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
-        
-        # Concatenate images side by side
-        w, h = z0_img_pil.width, z0_img_pil.height
-        combined_image = Image.new('RGB', (w*5, h))
-        combined_image.paste(z0_img_pil, (0, 0))
-        combined_image.paste(ns_img_pil, (w, 0))
-        combined_image.paste(zt_img_pil, (w*2, 0))
-        combined_image.paste(reg_img_pil, (w*3, 0))
-        combined_image.paste(org_img_pil, (w*4, 0))
-        combined_image.save(os.path.join(args.output_dir, f'visualization/vsd/{global_step}.png'))
-        #-------- Visualization --------#
-    
-    def diff_loss(self, z0, prompt_embeds, pooled_embeds, net_lpips, args):
-        
-        device = self.transformer_reg.device
-        u = compute_density_for_timestep_sampling(
-                        weighting_scheme="uniform",
-                        batch_size=1,
-                        logit_mean=0.0,
-                        logit_std=1.0,
-                        mode_scale=1.29,
-                    )
-        
-        t_idx = (u*1000).long().to(device)
-        self.model.scheduler.set_timesteps(1000, device=device)
-        times = self.model.scheduler.timesteps
-        t = times[t_idx]
-        sigma = t / 1000
-        
-        z0 = z0.to(device)
-        z0, prompt_embeds = z0.detach(), prompt_embeds.detach()
-        noise = torch.randn_like(z0)
-        sigma = sigma.half()
-        zt = (1-sigma) * z0 + sigma * noise         # noisy latents
-
-        # v-prediction
-        v_pred = self.predict_vector_reg(zt, t, prompt_embeds.to(device), pooled_embeds.to(device))
-        model_pred = v_pred * (-sigma) + zt
-        target = z0
-        
-        loss_weight = compute_loss_weighting_for_sd3("logit_normal", sigma)
-        diffusion_loss = loss_weight.float() * F.mse_loss(model_pred.float(), target.float()) 
-        
-        loss_d = diffusion_loss
-
-        return loss_d.mean()
-
-class OSEDiff_SD3_TEST(torch.nn.Module):
-    def __init__(self, args, base_model):
-        super().__init__()
-        
-        self.args = args
-        self.model = base_model
-        self.lora_path = args.lora_path
-        self.vae_path = args.vae_path
-        
-        # Add lora to transformer
-        print(f'Loading LoRA to Transformer from {self.lora_path}')
-        self.model.transformer.requires_grad_(False)
-        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
-        for name, param in self.model.transformer.named_parameters():
-            param.requires_grad = False
-        
-        # Insert LoRA into VAE
-        print(f"Loading LoRA to VAE from {self.vae_path}")
-        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
-        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu", weights_only=True)
-        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
-
-    def predict_vector(self, z, t, prompt_emb, pooled_emb):
-        v = self.model.transformer(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-
-    @torch.no_grad()
-    def forward(self, x_src, prompt):
-        
-        z_src = self.model.vae.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device)).latent_dist.sample() * self.model.vae.config.scaling_factor
-        
-        z_src = z_src.to(self.model.transformer.device)
-        
-        # calculate prompt_embeddings and neg_prompt_embeddings
-        batch_size, _, _, _ = x_src.shape
-        with torch.no_grad():
-            prompt_embeds, pooled_embeds = self.model.encode_prompt([prompt], batch_size)
-        
-        self.model.scheduler.set_timesteps(1, device=self.model.device)
-        timesteps = self.model.scheduler.timesteps
-
-        # Solve ODE
-        t = timesteps[0]
-        timestep = t.expand(z_src.shape[0]).to(self.model.transformer.device)
-        prompt_embeds = prompt_embeds.to(self.model.transformer.device, dtype=torch.float32)
-        pooled_embeds = pooled_embeds.to(self.model.transformer.device, dtype=torch.float32)
-        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
-        
-        z_src = z_src - pred_v
-        
-        with torch.no_grad():
-            output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
-
-        return output_image
-
-
-class OSEDiff_SD3_TEST_efficient(torch.nn.Module):
-    def __init__(self, args, base_model):
-        super().__init__()
-        
-        self.args = args
-        self.model = base_model
-        self.lora_path = args.lora_path
-        self.vae_path = args.vae_path
-        
-        # Add lora to transformer
-        print(f'Loading LoRA to Transformer from {self.lora_path}')
-        self.model.transformer.requires_grad_(False)
-        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
-        for name, param in self.model.transformer.named_parameters():
-            param.requires_grad = False
-        
-        # Insert LoRA into VAE
-        print(f"Loading LoRA to VAE from {self.vae_path}")
-        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
-        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu", weights_only=True)
-        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
-
-    def predict_vector(self, z, t, prompt_emb, pooled_emb):
-        v = self.model.transformer(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-
-    @torch.no_grad()
-    def forward(self, x_src, prompt):
-        
-        z_src = self.model.vae.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device)).latent_dist.sample() * self.model.vae.config.scaling_factor
-        
-        z_src = z_src.to(self.model.transformer.device)
-        
-        # calculate prompt_embeddings
-        batch_size, _, _, _ = x_src.shape
-        prompt_embeds, pooled_embeds = self.model.encode_prompt([prompt], batch_size)
-        
-        self.model.scheduler.set_timesteps(1, device=self.model.device)
-        timesteps = self.model.scheduler.timesteps
-
-        # Solve ODE
-        t = timesteps[0]
-        timestep = t.expand(z_src.shape[0]).to(self.model.transformer.device)
-        prompt_embeds = prompt_embeds.to(self.model.transformer.device, dtype=torch.float32)
-        pooled_embeds = pooled_embeds.to(self.model.transformer.device, dtype=torch.float32)
-        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
-        z_src = z_src - pred_v
-        
-        output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
-
-        return output_image
-
-
-class OSEDiff_SD3_TEST_TILE(torch.nn.Module):
-    def __init__(self, args, base_model):
-        super().__init__()
-        
-        self.args = args
-        self.model = base_model
-        self.lora_path = args.lora_path
-        self.vae_path = args.vae_path
-        
-        # Add lora to transformer
-        print(f'Loading LoRA to Transformer from {self.lora_path}')
-        self.model.transformer.requires_grad_(False)
-        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
-        for name, param in self.model.transformer.named_parameters():
-            param.requires_grad = False
-        
-        # Insert LoRA into VAE
-        print(f"Loading LoRA to VAE from {self.vae_path}")
-        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
-        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu", weights_only=True)
-        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
-        
-        # save original forward (only once)
-        if not hasattr(self.model.vae.encoder, 'original_forward'):
-            setattr(self.model.vae.encoder, 'original_forward', self.model.vae.encoder.forward)
-        if not hasattr(self.model.vae.decoder, 'original_forward'):
-            setattr(self.model.vae.decoder, 'original_forward', self.model.vae.decoder.forward)
-        encoder_tile = args.vae_encoder_tiled_size
-        decoder_tile = args.vae_decoder_tiled_size
-        self.model.vae.encoder.forward = VAEHook(
-            self.model.vae.encoder,
-            tile_size=encoder_tile,
-            is_decoder=False,
-            fast_decoder=False,
-            fast_encoder=True,
-            color_fix=False,
-            to_gpu=True
-        )
-        self.model.vae.decoder.forward = VAEHook(
-            self.model.vae.decoder,
-            tile_size=decoder_tile,
-            is_decoder=True,
-            fast_decoder=True,
-            fast_encoder=False,
-            color_fix=False,
-            to_gpu=True
-        )
-
-    def predict_vector(self, z, t, prompt_emb, pooled_emb):
-        v = self.model.transformer(hidden_states=z,
-                             timestep=t,
-                             pooled_projections=pooled_emb,
-                             encoder_hidden_states=prompt_emb,
-                             return_dict=False)[0]
-        return v
-    
-    @torch.no_grad()
-    def create_full_latent(self, x_full: torch.Tensor, vlm_model, vlm_processor, full_path, next_path, prompt_type) -> torch.Tensor:
-        device = self.model.transformer.device
-        # 1) encode to full latent (via VAEHook)
-        z_full = self.model.vae.encode(x_full).latent_dist.sample() \
-                * self.model.vae.config.scaling_factor
-        z_full = z_full.to(device)
-        B, C, H, W = z_full.shape
-
-        # 2) grid size
-        tsize = self.args.latent_tiled_size
-        tover  = self.args.latent_tiled_overlap
-        stride = tsize - tover
-        rows = (H - tsize + stride - 1)//stride + 1
-        cols = (W - tsize + stride - 1)//stride + 1
-        print(f'TILE SIZE: {tsize}, TILE OVERLAP: {tover}, STRIDE: {stride}, ROWS: {rows}, COLS: {cols}')
-
-        # 3) make gaussian weight patched [B,C,tsize,tsize]
-        weights = self._make_gaussian(tsize, tsize, 1).to(z_full.device)
-
-        # 4) collect all patches
-        positions = []
-        out_tiles = []
-        timestep = self.model.scheduler.timesteps[0].expand(B).to(z_full.device)
-
-        for i in range(rows):
-            y0 = min(i*stride, H - tsize)
-            for j in range(cols):
-                x0 = min(j*stride, W - tsize)
-                positions.append((y0, x0))
-                patch = z_full[:, :, y0:y0+tsize, x0:x0+tsize]
-                
-                # decode and save patch for later usage
-                patch_path = f'{next_path[:-4]}_patch_row{i}col{j}.png'
-                patch_img = self.decode_full_latent(patch)
-                patch_pil = transforms.ToPILImage()((patch_img[0] * 0.5 + 0.5).clamp(0,1))
-                patch_pil.save(patch_path)
-                
-                # create prompt to explain patch (CoZ)
-                prompt = self.create_prompt(vlm_model, vlm_processor, full_path, patch_path, prompt_type)
-                print('PROMPT: ', prompt)
-                prompt_emb, pooled_emb = self.model.encode_prompt([prompt], batch_size=B)
-                prompt_emb = prompt_emb.to(z_full.device, dtype=torch.float32)
-                pooled_emb = pooled_emb.to(z_full.device, dtype=torch.float32)
-                
-                v = self.predict_vector(patch, timestep, prompt_emb, pooled_emb)
-                out_tiles.append(patch - v)
-
-        # 5) accumulate + normalize
-        z_out = torch.zeros_like(z_full)
-        z_norm = torch.zeros_like(z_full)
-        norm  = torch.zeros_like(z_full)
-
-        for (y0,x0), tile in zip(positions, out_tiles):
-            z_out[:, :, y0:y0+tsize, x0:x0+tsize] += tile
-            z_norm[:, :, y0:y0+tsize, x0:x0+tsize] += tile * weights
-            norm[:, :, y0:y0+tsize, x0:x0+tsize] += weights
-
-        # 6) avoid division by zero and finalize
-        eps = 1e-10
-        z_norm = z_norm / (norm + eps)
-        return z_norm, z_out
-
-    @torch.no_grad()
-    def decode_full_latent(self, z_full: torch.Tensor) -> torch.Tensor:
-        """
-        Decode the tiled full latent into an RGB image (with tiled VAE decoder).
-        """
-        z_full = z_full.to(self.model.vae.device)
-        img = self.model.vae.decode(z_full / self.model.vae.config.scaling_factor).sample
-        return img.clamp(-1,1)
-    
-    @torch.no_grad()
-    def create_prompt(self, vlm_model, vlm_processor, full_path, patch_path, prompt_type):
-        if prompt_type in ('vlm','vlm_base'):
-            from qwen_vl_utils import process_vision_info
-            
-            message_text = None
-            start_image_path = full_path
-            input_image_path = patch_path
-            
-            message_text = "The second image is a zoom-in of the first image. Based on this knowledge, what is in the second image? Give me a set of words."
-            messages = [
-                {"role": "system", "content": f"{message_text}"},
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "image", "image": start_image_path},
-                        {"type": "image", "image": input_image_path}
-                    ]
-                }
-            ]
-
-            text = vlm_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-            image_inputs, video_inputs = process_vision_info(messages)
-            inputs = vlm_processor(
-                text=[text],
-                images=image_inputs,
-                videos=video_inputs,
-                padding=True,
-                return_tensors="pt",
-            )
-            generated_ids = vlm_model.generate(**inputs, max_new_tokens=16)
-            generated_ids_trimmed = [
-                out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-            ]
-            output_text = vlm_processor.batch_decode(
-                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-            )
-
-            prompt_text = output_text[0]
-            return prompt_text
-        else:
-            raise ValueError(f"Unknown prompt_type: {prompt_type}")
-    
-    def _make_gaussian(self, w, h, nb):
-        from numpy import pi, exp, sqrt
-        import numpy as np
-
-        latent_width = w
-        latent_height = h
-
-        var = 0.01
-        midpoint = (latent_width - 1) / 2       # -1 because index goes from 0 to latent_width - 1
-        x_probs = [exp(-(x-midpoint)*(x-midpoint)/(latent_width*latent_width)/(2*var)) / sqrt(2*pi*var) for x in range(latent_width)]
-        midpoint = latent_height / 2
-        y_probs = [exp(-(y-midpoint)*(y-midpoint)/(latent_height*latent_height)/(2*var)) / sqrt(2*pi*var) for y in range(latent_height)]
-
-        weights = np.outer(y_probs, x_probs)
-        return torch.tile(torch.tensor(weights, device=self.model.vae.device), (nb, self.model.vae.config.latent_channels, 1, 1))

vendor/utils/devices.py

@@ -1,138 +0,0 @@
-import sys
-import contextlib
-from functools import lru_cache
-
-import torch
-#from modules import errors
-
-if sys.platform == "darwin":
-    from modules import mac_specific
-
-
-def has_mps() -> bool:
-    if sys.platform != "darwin":
-        return False
-    else:
-        return mac_specific.has_mps
-
-
-def get_cuda_device_string():
-    return "cuda"
-
-
-def get_optimal_device_name():
-    if torch.cuda.is_available():
-        return get_cuda_device_string()
-
-    if has_mps():
-        return "mps"
-
-    return "cpu"
-
-
-def get_optimal_device():
-    return torch.device(get_optimal_device_name())
-
-
-def get_device_for(task):
-    return get_optimal_device()
-
-
-def torch_gc():
-
-    if torch.cuda.is_available():
-        with torch.cuda.device(get_cuda_device_string()):
-            torch.cuda.empty_cache()
-            torch.cuda.ipc_collect()
-
-    if has_mps():
-        mac_specific.torch_mps_gc()
-
-
-def enable_tf32():
-    if torch.cuda.is_available():
-
-        # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
-        # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
-        if any(torch.cuda.get_device_capability(devid) == (7, 5) for devid in range(0, torch.cuda.device_count())):
-            torch.backends.cudnn.benchmark = True
-
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.backends.cudnn.allow_tf32 = True
-
-
-enable_tf32()
-#errors.run(enable_tf32, "Enabling TF32")
-
-cpu = torch.device("cpu")
-device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = torch.device("cuda")
-dtype = torch.float16
-dtype_vae = torch.float16
-dtype_unet = torch.float16
-unet_needs_upcast = False
-
-
-def cond_cast_unet(input):
-    return input.to(dtype_unet) if unet_needs_upcast else input
-
-
-def cond_cast_float(input):
-    return input.float() if unet_needs_upcast else input
-
-
-def randn(seed, shape):
-    torch.manual_seed(seed)
-    return torch.randn(shape, device=device)
-
-
-def randn_without_seed(shape):
-    return torch.randn(shape, device=device)
-
-
-def autocast(disable=False):
-    if disable:
-        return contextlib.nullcontext()
-
-    return torch.autocast("cuda")
-
-
-def without_autocast(disable=False):
-    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
-
-
-class NansException(Exception):
-    pass
-
-
-def test_for_nans(x, where):
-    if not torch.all(torch.isnan(x)).item():
-        return
-
-    if where == "unet":
-        message = "A tensor with all NaNs was produced in Unet."
-
-    elif where == "vae":
-        message = "A tensor with all NaNs was produced in VAE."
-
-    else:
-        message = "A tensor with all NaNs was produced."
-
-    message += " Use --disable-nan-check commandline argument to disable this check."
-
-    raise NansException(message)
-
-
-@lru_cache
-def first_time_calculation():
-    """
-    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
-    spends about 2.7 seconds doing that, at least wih NVidia.
-    """
-
-    x = torch.zeros((1, 1)).to(device, dtype)
-    linear = torch.nn.Linear(1, 1).to(device, dtype)
-    linear(x)
-
-    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
-    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
-    conv2d(x)

vendor/utils/vaehook.py

@@ -1,829 +0,0 @@
-# ------------------------------------------------------------------------
-#
-#   Ultimate VAE Tile Optimization
-#
-#   Introducing a revolutionary new optimization designed to make
-#   the VAE work with giant images on limited VRAM!
-#   Say goodbye to the frustration of OOM and hello to seamless output!
-#
-# ------------------------------------------------------------------------
-#
-#   This script is a wild hack that splits the image into tiles,
-#   encodes each tile separately, and merges the result back together.
-#
-#   Advantages:
-#   - The VAE can now work with giant images on limited VRAM
-#       (~10 GB for 8K images!)
-#   - The merged output is completely seamless without any post-processing.
-#
-#   Drawbacks:
-#   - Giant RAM needed. To store the intermediate results for a 4096x4096
-#       images, you need 32 GB RAM it consumes ~20GB); for 8192x8192
-#       you need 128 GB RAM machine (it consumes ~100 GB)
-#   - NaNs always appear in for 8k images when you use fp16 (half) VAE
-#       You must use --no-half-vae to disable half VAE for that giant image.
-#   - Slow speed. With default tile size, it takes around 50/200 seconds
-#       to encode/decode a 4096x4096 image; and 200/900 seconds to encode/decode
-#       a 8192x8192 image. (The speed is limited by both the GPU and the CPU.)
-#   - The gradient calculation is not compatible with this hack. It
-#       will break any backward() or torch.autograd.grad() that passes VAE.
-#       (But you can still use the VAE to generate training data.)
-#
-#   How it works:
-#   1) The image is split into tiles.
-#       - To ensure perfect results, each tile is padded with 32 pixels
-#           on each side.
-#       - Then the conv2d/silu/upsample/downsample can produce identical
-#           results to the original image without splitting.
-#   2) The original forward is decomposed into a task queue and a task worker.
-#       - The task queue is a list of functions that will be executed in order.
-#       - The task worker is a loop that executes the tasks in the queue.
-#   3) The task queue is executed for each tile.
-#       - Current tile is sent to GPU.
-#       - local operations are directly executed.
-#       - Group norm calculation is temporarily suspended until the mean
-#           and var of all tiles are calculated.
-#       - The residual is pre-calculated and stored and addded back later.
-#       - When need to go to the next tile, the current tile is send to cpu.
-#   4) After all tiles are processed, tiles are merged on cpu and return.
-#
-#   Enjoy!
-#
-#   @author: LI YI @ Nanyang Technological University - Singapore
-#   @date: 2023-03-02
-#   @license: MIT License
-#
-#   Please give me a star if you like this project!
-#
-# -------------------------------------------------------------------------
-
-import gc
-from time import time
-import math
-from tqdm import tqdm
-
-import torch
-import torch.version
-import torch.nn.functional as F
-from einops import rearrange
-import os
-import sys
-sys.path.append(os.getcwd())
-import utils.devices as devices
-
-try:
-    import xformers
-    import xformers.ops
-except ImportError:
-    pass
-
-sd_flag = False
-
-def get_recommend_encoder_tile_size():
-    if torch.cuda.is_available():
-        total_memory = torch.cuda.get_device_properties(
-            devices.device).total_memory // 2**20
-        if total_memory > 16*1000:
-            ENCODER_TILE_SIZE = 3072
-        elif total_memory > 12*1000:
-            ENCODER_TILE_SIZE = 2048
-        elif total_memory > 8*1000:
-            ENCODER_TILE_SIZE = 1536
-        else:
-            ENCODER_TILE_SIZE = 960
-    else:
-        ENCODER_TILE_SIZE = 512
-    return ENCODER_TILE_SIZE
-
-
-def get_recommend_decoder_tile_size():
-    if torch.cuda.is_available():
-        total_memory = torch.cuda.get_device_properties(
-            devices.device).total_memory // 2**20
-        if total_memory > 30*1000:
-            DECODER_TILE_SIZE = 256
-        elif total_memory > 16*1000:
-            DECODER_TILE_SIZE = 192
-        elif total_memory > 12*1000:
-            DECODER_TILE_SIZE = 128
-        elif total_memory > 8*1000:
-            DECODER_TILE_SIZE = 96
-        else:
-            DECODER_TILE_SIZE = 64
-    else:
-        DECODER_TILE_SIZE = 64
-    return DECODER_TILE_SIZE
-
-
-if 'global const':
-    DEFAULT_ENABLED = False
-    DEFAULT_MOVE_TO_GPU = False
-    DEFAULT_FAST_ENCODER = True
-    DEFAULT_FAST_DECODER = True
-    DEFAULT_COLOR_FIX = 0
-    DEFAULT_ENCODER_TILE_SIZE = get_recommend_encoder_tile_size()
-    DEFAULT_DECODER_TILE_SIZE = get_recommend_decoder_tile_size()
-
-
-# inplace version of silu
-def inplace_nonlinearity(x):
-    # Test: fix for Nans
-    return F.silu(x, inplace=True)
-
-# extracted from ldm.modules.diffusionmodules.model
-
-# from diffusers lib
-def attn_forward_new(self, h_):
-    batch_size, channel, height, width = h_.shape
-    hidden_states = h_.view(batch_size, channel, height * width).transpose(1, 2)
-
-    attention_mask = None
-    encoder_hidden_states = None
-    batch_size, sequence_length, _ = hidden_states.shape
-    attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-
-    query = self.to_q(hidden_states)
-
-    if encoder_hidden_states is None:
-        encoder_hidden_states = hidden_states
-    elif self.norm_cross:
-        encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
-
-    key = self.to_k(encoder_hidden_states)
-    value = self.to_v(encoder_hidden_states)
-
-    query = self.head_to_batch_dim(query)
-    key = self.head_to_batch_dim(key)
-    value = self.head_to_batch_dim(value)
-
-    attention_probs = self.get_attention_scores(query, key, attention_mask)
-    hidden_states = torch.bmm(attention_probs, value)
-    hidden_states = self.batch_to_head_dim(hidden_states)
-
-    # linear proj
-    hidden_states = self.to_out[0](hidden_states)
-    # dropout
-    hidden_states = self.to_out[1](hidden_states)
-
-    hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-    return hidden_states
-
-def attn_forward(self, h_):
-    q = self.q(h_)
-    k = self.k(h_)
-    v = self.v(h_)
-
-    # compute attention
-    b, c, h, w = q.shape
-    q = q.reshape(b, c, h*w)
-    q = q.permute(0, 2, 1)   # b,hw,c
-    k = k.reshape(b, c, h*w)  # b,c,hw
-    w_ = torch.bmm(q, k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-    w_ = w_ * (int(c)**(-0.5))
-    w_ = torch.nn.functional.softmax(w_, dim=2)
-
-    # attend to values
-    v = v.reshape(b, c, h*w)
-    w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
-    # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-    h_ = torch.bmm(v, w_)
-    h_ = h_.reshape(b, c, h, w)
-
-    h_ = self.proj_out(h_)
-
-    return h_
-
-
-def xformer_attn_forward(self, h_):
-    q = self.q(h_)
-    k = self.k(h_)
-    v = self.v(h_)
-
-    # compute attention
-    B, C, H, W = q.shape
-    q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
-
-    q, k, v = map(
-        lambda t: t.unsqueeze(3)
-        .reshape(B, t.shape[1], 1, C)
-        .permute(0, 2, 1, 3)
-        .reshape(B * 1, t.shape[1], C)
-        .contiguous(),
-        (q, k, v),
-    )
-    out = xformers.ops.memory_efficient_attention(
-        q, k, v, attn_bias=None, op=self.attention_op)
-
-    out = (
-        out.unsqueeze(0)
-        .reshape(B, 1, out.shape[1], C)
-        .permute(0, 2, 1, 3)
-        .reshape(B, out.shape[1], C)
-    )
-    out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
-    out = self.proj_out(out)
-    return out
-
-
-def attn2task(task_queue, net):
-    if False: #isinstance(net, AttnBlock):
-        task_queue.append(('store_res', lambda x: x))
-        task_queue.append(('pre_norm', net.norm))
-        task_queue.append(('attn', lambda x, net=net: attn_forward(net, x)))
-        task_queue.append(['add_res', None])
-    elif False: #isinstance(net, MemoryEfficientAttnBlock):
-        task_queue.append(('store_res', lambda x: x))
-        task_queue.append(('pre_norm', net.norm))
-        task_queue.append(
-            ('attn', lambda x, net=net: xformer_attn_forward(net, x)))
-        task_queue.append(['add_res', None])
-    else:
-        task_queue.append(('store_res', lambda x: x))
-        task_queue.append(('pre_norm', net.group_norm))
-        task_queue.append(('attn', lambda x, net=net: attn_forward_new(net, x)))
-        task_queue.append(['add_res', None])
-
-def resblock2task(queue, block):
-    """
-    Turn a ResNetBlock into a sequence of tasks and append to the task queue
-
-    @param queue: the target task queue
-    @param block: ResNetBlock
-
-    """
-    if block.in_channels != block.out_channels:
-        if sd_flag:
-            if block.use_conv_shortcut:
-                queue.append(('store_res', block.conv_shortcut))
-            else:
-                queue.append(('store_res', block.nin_shortcut))
-        else:
-            if block.use_in_shortcut:
-                queue.append(('store_res', block.conv_shortcut))
-            else:
-                queue.append(('store_res', block.nin_shortcut))
-
-    else:
-        queue.append(('store_res', lambda x: x))
-    queue.append(('pre_norm', block.norm1))
-    queue.append(('silu', inplace_nonlinearity))
-    queue.append(('conv1', block.conv1))
-    queue.append(('pre_norm', block.norm2))
-    queue.append(('silu', inplace_nonlinearity))
-    queue.append(('conv2', block.conv2))
-    queue.append(['add_res', None])
-
-
-
-def build_sampling(task_queue, net, is_decoder):
-    """
-    Build the sampling part of a task queue
-    @param task_queue: the target task queue
-    @param net: the network
-    @param is_decoder: currently building decoder or encoder
-    """
-    if is_decoder:
-        # resblock2task(task_queue, net.mid.block_1)
-        # attn2task(task_queue, net.mid.attn_1)
-        # resblock2task(task_queue, net.mid.block_2)
-        # resolution_iter = reversed(range(net.num_resolutions))
-        # block_ids = net.num_res_blocks + 1
-        # condition = 0
-        # module = net.up
-        # func_name = 'upsample'
-        resblock2task(task_queue, net.mid_block.resnets[0])
-        attn2task(task_queue, net.mid_block.attentions[0])
-        resblock2task(task_queue, net.mid_block.resnets[1])
-        resolution_iter = (range(len(net.up_blocks)))  # range(0,4)
-        block_ids = 2 + 1
-        condition = len(net.up_blocks) - 1
-        module = net.up_blocks
-        func_name = 'upsamplers'
-    else:
-        # resolution_iter = range(net.num_resolutions)
-        # block_ids = net.num_res_blocks
-        # condition = net.num_resolutions - 1
-        # module = net.down
-        # func_name = 'downsample'
-        resolution_iter = (range(len(net.down_blocks)))  # range(0,4)
-        block_ids = 2
-        condition = len(net.down_blocks) - 1
-        module = net.down_blocks
-        func_name = 'downsamplers'
-
-
-    for i_level in resolution_iter:
-        for i_block in range(block_ids):
-            resblock2task(task_queue, module[i_level].resnets[i_block])
-        if i_level != condition:
-            if is_decoder:
-                task_queue.append((func_name, module[i_level].upsamplers[0]))
-            else:
-                task_queue.append((func_name, module[i_level].downsamplers[0]))
-
-    if not is_decoder:
-        resblock2task(task_queue, net.mid_block.resnets[0])
-        attn2task(task_queue, net.mid_block.attentions[0])
-        resblock2task(task_queue, net.mid_block.resnets[1])
-
-
-def build_task_queue(net, is_decoder):
-    """
-    Build a single task queue for the encoder or decoder
-    @param net: the VAE decoder or encoder network
-    @param is_decoder: currently building decoder or encoder
-    @return: the task queue
-    """
-    task_queue = []
-    task_queue.append(('conv_in', net.conv_in))
-
-    # construct the sampling part of the task queue
-    # because encoder and decoder share the same architecture, we extract the sampling part
-    build_sampling(task_queue, net, is_decoder)
-    if is_decoder and not sd_flag:
-        net.give_pre_end = False
-        net.tanh_out = False
-
-    if not is_decoder or not net.give_pre_end:
-        if sd_flag:
-            task_queue.append(('pre_norm', net.norm_out))
-        else:
-            task_queue.append(('pre_norm', net.conv_norm_out))
-        task_queue.append(('silu', inplace_nonlinearity))
-        task_queue.append(('conv_out', net.conv_out))
-        if is_decoder and net.tanh_out:
-            task_queue.append(('tanh', torch.tanh))
-
-    return task_queue
-
-
-def clone_task_queue(task_queue):
-    """
-    Clone a task queue
-    @param task_queue: the task queue to be cloned
-    @return: the cloned task queue
-    """
-    return [[item for item in task] for task in task_queue]
-
-
-def get_var_mean(input, num_groups, eps=1e-6):
-    """
-    Get mean and var for group norm
-    """
-    b, c = input.size(0), input.size(1)
-    channel_in_group = int(c/num_groups)
-    input_reshaped = input.contiguous().view(
-        1, int(b * num_groups), channel_in_group, *input.size()[2:])
-    var, mean = torch.var_mean(
-        input_reshaped, dim=[0, 2, 3, 4], unbiased=False)
-    return var, mean
-
-
-def custom_group_norm(input, num_groups, mean, var, weight=None, bias=None, eps=1e-6):
-    """
-    Custom group norm with fixed mean and var
-
-    @param input: input tensor
-    @param num_groups: number of groups. by default, num_groups = 32
-    @param mean: mean, must be pre-calculated by get_var_mean
-    @param var: var, must be pre-calculated by get_var_mean
-    @param weight: weight, should be fetched from the original group norm
-    @param bias: bias, should be fetched from the original group norm
-    @param eps: epsilon, by default, eps = 1e-6 to match the original group norm
-
-    @return: normalized tensor
-    """
-    b, c = input.size(0), input.size(1)
-    channel_in_group = int(c/num_groups)
-    input_reshaped = input.contiguous().view(
-        1, int(b * num_groups), channel_in_group, *input.size()[2:])
-
-    out = F.batch_norm(input_reshaped, mean, var, weight=None, bias=None,
-                       training=False, momentum=0, eps=eps)
-
-    out = out.view(b, c, *input.size()[2:])
-
-    # post affine transform
-    if weight is not None:
-        out *= weight.view(1, -1, 1, 1)
-    if bias is not None:
-        out += bias.view(1, -1, 1, 1)
-    return out
-
-
-def crop_valid_region(x, input_bbox, target_bbox, is_decoder):
-    """
-    Crop the valid region from the tile
-    @param x: input tile
-    @param input_bbox: original input bounding box
-    @param target_bbox: output bounding box
-    @param scale: scale factor
-    @return: cropped tile
-    """
-    padded_bbox = [i * 8 if is_decoder else i//8 for i in input_bbox]
-    margin = [target_bbox[i] - padded_bbox[i] for i in range(4)]
-    return x[:, :, margin[2]:x.size(2)+margin[3], margin[0]:x.size(3)+margin[1]]
-
-# ↓↓↓ https://github.com/Kahsolt/stable-diffusion-webui-vae-tile-infer ↓↓↓
-
-
-def perfcount(fn):
-    def wrapper(*args, **kwargs):
-        ts = time()
-
-        if torch.cuda.is_available():
-            torch.cuda.reset_peak_memory_stats(devices.device)
-        devices.torch_gc()
-        gc.collect()
-
-        ret = fn(*args, **kwargs)
-
-        devices.torch_gc()
-        gc.collect()
-        if torch.cuda.is_available():
-            vram = torch.cuda.max_memory_allocated(devices.device) / 2**20
-            torch.cuda.reset_peak_memory_stats(devices.device)
-            print(
-                f'[Tiled VAE]: Done in {time() - ts:.3f}s, max VRAM alloc {vram:.3f} MB')
-        else:
-            print(f'[Tiled VAE]: Done in {time() - ts:.3f}s')
-
-        return ret
-    return wrapper
-
-# copy end :)
-
-
-class GroupNormParam:
-    def __init__(self):
-        self.var_list = []
-        self.mean_list = []
-        self.pixel_list = []
-        self.weight = None
-        self.bias = None
-
-    def add_tile(self, tile, layer):
-        var, mean = get_var_mean(tile, 32)
-        # For giant images, the variance can be larger than max float16
-        # In this case we create a copy to float32
-        if var.dtype == torch.float16 and var.isinf().any():
-            fp32_tile = tile.float()
-            var, mean = get_var_mean(fp32_tile, 32)
-        # ============= DEBUG: test for infinite =============
-        # if torch.isinf(var).any():
-        #    print('var: ', var)
-        # ====================================================
-        self.var_list.append(var)
-        self.mean_list.append(mean)
-        self.pixel_list.append(
-            tile.shape[2]*tile.shape[3])
-        if hasattr(layer, 'weight'):
-            self.weight = layer.weight
-            self.bias = layer.bias
-        else:
-            self.weight = None
-            self.bias = None
-
-    def summary(self):
-        """
-        summarize the mean and var and return a function
-        that apply group norm on each tile
-        """
-        if len(self.var_list) == 0:
-            return None
-        var = torch.vstack(self.var_list)
-        mean = torch.vstack(self.mean_list)
-        max_value = max(self.pixel_list)
-        pixels = torch.tensor(
-            self.pixel_list, dtype=torch.float32, device=devices.device) / max_value
-        sum_pixels = torch.sum(pixels)
-        pixels = pixels.unsqueeze(
-            1) / sum_pixels
-        var = torch.sum(
-            var * pixels, dim=0)
-        mean = torch.sum(
-            mean * pixels, dim=0)
-        return lambda x:  custom_group_norm(x, 32, mean, var, self.weight, self.bias)
-
-    @staticmethod
-    def from_tile(tile, norm):
-        """
-        create a function from a single tile without summary
-        """
-        var, mean = get_var_mean(tile, 32)
-        if var.dtype == torch.float16 and var.isinf().any():
-            fp32_tile = tile.float()
-            var, mean = get_var_mean(fp32_tile, 32)
-            # if it is a macbook, we need to convert back to float16
-            if var.device.type == 'mps':
-                # clamp to avoid overflow
-                var = torch.clamp(var, 0, 60000)
-                var = var.half()
-                mean = mean.half()
-        if hasattr(norm, 'weight'):
-            weight = norm.weight
-            bias = norm.bias
-        else:
-            weight = None
-            bias = None
-
-        def group_norm_func(x, mean=mean, var=var, weight=weight, bias=bias):
-            return custom_group_norm(x, 32, mean, var, weight, bias, 1e-6)
-        return group_norm_func
-
-
-class VAEHook:
-    def __init__(self, net, tile_size, is_decoder, fast_decoder, fast_encoder, color_fix, to_gpu=False):
-        self.net = net                  # encoder | decoder
-        self.tile_size = tile_size
-        self.is_decoder = is_decoder
-        self.fast_mode = (fast_encoder and not is_decoder) or (
-            fast_decoder and is_decoder)
-        self.color_fix = color_fix and not is_decoder
-        self.to_gpu = to_gpu
-        self.pad = 11 if is_decoder else 32
-
-    def __call__(self, x):
-        B, C, H, W = x.shape
-        original_device = next(self.net.parameters()).device
-        try:
-            if self.to_gpu:
-                # self.net.to(devices.get_optimal_device())
-                self.net.to(original_device)
-            if max(H, W) <= self.pad * 2 + self.tile_size:
-                # print("[Tiled VAE]: the input size is tiny and unnecessary to tile.")
-                return self.net.original_forward(x).to(original_device)
-            else:
-                return self.vae_tile_forward(x)
-        finally:
-            self.net.to(original_device)
-
-    def get_best_tile_size(self, lowerbound, upperbound):
-        """
-        Get the best tile size for GPU memory
-        """
-        divider = 32
-        while divider >= 2:
-            remainer = lowerbound % divider
-            if remainer == 0:
-                return lowerbound
-            candidate = lowerbound - remainer + divider
-            if candidate <= upperbound:
-                return candidate
-            divider //= 2
-        return lowerbound
-
-    def split_tiles(self, h, w):
-        """
-        Tool function to split the image into tiles
-        @param h: height of the image
-        @param w: width of the image
-        @return: tile_input_bboxes, tile_output_bboxes
-        """
-        tile_input_bboxes, tile_output_bboxes = [], []
-        tile_size = self.tile_size
-        pad = self.pad
-        num_height_tiles = math.ceil((h - 2 * pad) / tile_size)
-        num_width_tiles = math.ceil((w - 2 * pad) / tile_size)
-        # If any of the numbers are 0, we let it be 1
-        # This is to deal with long and thin images
-        num_height_tiles = max(num_height_tiles, 1)
-        num_width_tiles = max(num_width_tiles, 1)
-
-        # Suggestions from https://github.com/Kahsolt: auto shrink the tile size
-        real_tile_height = math.ceil((h - 2 * pad) / num_height_tiles)
-        real_tile_width = math.ceil((w - 2 * pad) / num_width_tiles)
-        real_tile_height = self.get_best_tile_size(real_tile_height, tile_size)
-        real_tile_width = self.get_best_tile_size(real_tile_width, tile_size)
-
-        print(f'[Tiled VAE]: split to {num_height_tiles}x{num_width_tiles} = {num_height_tiles*num_width_tiles} tiles. ' +
-              f'Optimal tile size {real_tile_width}x{real_tile_height}, original tile size {tile_size}x{tile_size}')
-
-        for i in range(num_height_tiles):
-            for j in range(num_width_tiles):
-                # bbox: [x1, x2, y1, y2]
-                # the padding is is unnessary for image borders. So we directly start from (32, 32)
-                input_bbox = [
-                    pad + j * real_tile_width,
-                    min(pad + (j + 1) * real_tile_width, w),
-                    pad + i * real_tile_height,
-                    min(pad + (i + 1) * real_tile_height, h),
-                ]
-
-                # if the output bbox is close to the image boundary, we extend it to the image boundary
-                output_bbox = [
-                    input_bbox[0] if input_bbox[0] > pad else 0,
-                    input_bbox[1] if input_bbox[1] < w - pad else w,
-                    input_bbox[2] if input_bbox[2] > pad else 0,
-                    input_bbox[3] if input_bbox[3] < h - pad else h,
-                ]
-
-                # scale to get the final output bbox
-                output_bbox = [x * 8 if self.is_decoder else x // 8 for x in output_bbox]
-                tile_output_bboxes.append(output_bbox)
-
-                # indistinguishable expand the input bbox by pad pixels
-                tile_input_bboxes.append([
-                    max(0, input_bbox[0] - pad),
-                    min(w, input_bbox[1] + pad),
-                    max(0, input_bbox[2] - pad),
-                    min(h, input_bbox[3] + pad),
-                ])
-
-        return tile_input_bboxes, tile_output_bboxes
-
-    @torch.no_grad()
-    def estimate_group_norm(self, z, task_queue, color_fix):
-        device = z.device
-        tile = z
-        last_id = len(task_queue) - 1
-        while last_id >= 0 and task_queue[last_id][0] != 'pre_norm':
-            last_id -= 1
-        if last_id <= 0 or task_queue[last_id][0] != 'pre_norm':
-            raise ValueError('No group norm found in the task queue')
-        # estimate until the last group norm
-        for i in range(last_id + 1):
-            task = task_queue[i]
-            if task[0] == 'pre_norm':
-                group_norm_func = GroupNormParam.from_tile(tile, task[1])
-                task_queue[i] = ('apply_norm', group_norm_func)
-                if i == last_id:
-                    return True
-                tile = group_norm_func(tile)
-            elif task[0] == 'store_res':
-                task_id = i + 1
-                while task_id < last_id and task_queue[task_id][0] != 'add_res':
-                    task_id += 1
-                if task_id >= last_id:
-                    continue
-                task_queue[task_id][1] = task[1](tile)
-            elif task[0] == 'add_res':
-                tile += task[1].to(device)
-                task[1] = None
-            elif color_fix and task[0] == 'downsample':
-                for j in range(i, last_id + 1):
-                    if task_queue[j][0] == 'store_res':
-                        task_queue[j] = ('store_res_cpu', task_queue[j][1])
-                return True
-            else:
-                tile = task[1](tile)
-            try:
-                devices.test_for_nans(tile, "vae")
-            except:
-                print(f'Nan detected in fast mode estimation. Fast mode disabled.')
-                return False
-
-        raise IndexError('Should not reach here')
-
-    # @perfcount
-    @torch.no_grad()
-    def vae_tile_forward(self, z):
-        """
-        Decode a latent vector z into an image in a tiled manner.
-        @param z: latent vector
-        @return: image
-        """
-        device = next(self.net.parameters()).device
-        net = self.net
-        tile_size = self.tile_size
-        is_decoder = self.is_decoder
-
-        z = z.detach() # detach the input to avoid backprop
-
-        N, height, width = z.shape[0], z.shape[2], z.shape[3]
-        net.last_z_shape = z.shape
-
-        # Split the input into tiles and build a task queue for each tile
-        print(f'[Tiled VAE]: input_size: {z.shape}, tile_size: {tile_size}, padding: {self.pad}')
-
-        in_bboxes, out_bboxes = self.split_tiles(height, width)
-
-        # Prepare tiles by split the input latents
-        tiles = []
-        for input_bbox in in_bboxes:
-            tile = z[:, :, input_bbox[2]:input_bbox[3], input_bbox[0]:input_bbox[1]].cpu()
-            tiles.append(tile)
-
-        num_tiles = len(tiles)
-        num_completed = 0
-
-        # Build task queues
-        single_task_queue = build_task_queue(net, is_decoder)
-        #print(single_task_queue)
-        if self.fast_mode:
-            # Fast mode: downsample the input image to the tile size,
-            # then estimate the group norm parameters on the downsampled image
-            scale_factor = tile_size / max(height, width)
-            z = z.to(device)
-            downsampled_z = F.interpolate(z, scale_factor=scale_factor, mode='nearest-exact')
-            # use nearest-exact to keep statictics as close as possible
-            print(f'[Tiled VAE]: Fast mode enabled, estimating group norm parameters on {downsampled_z.shape[3]} x {downsampled_z.shape[2]} image')
-
-            # ======= Special thanks to @Kahsolt for distribution shift issue ======= #
-            # The downsampling will heavily distort its mean and std, so we need to recover it.
-            std_old, mean_old = torch.std_mean(z, dim=[0, 2, 3], keepdim=True)
-            std_new, mean_new = torch.std_mean(downsampled_z, dim=[0, 2, 3], keepdim=True)
-            downsampled_z = (downsampled_z - mean_new) / std_new * std_old + mean_old
-            del std_old, mean_old, std_new, mean_new
-            # occasionally the std_new is too small or too large, which exceeds the range of float16
-            # so we need to clamp it to max z's range.
-            downsampled_z = torch.clamp_(downsampled_z, min=z.min(), max=z.max())
-            estimate_task_queue = clone_task_queue(single_task_queue)
-            if self.estimate_group_norm(downsampled_z, estimate_task_queue, color_fix=self.color_fix):
-                single_task_queue = estimate_task_queue
-            del downsampled_z
-
-        task_queues = [clone_task_queue(single_task_queue) for _ in range(num_tiles)]
-
-        # Dummy result
-        result = None
-        result_approx = None
-        #try:
-        #    with devices.autocast():
-        #        result_approx = torch.cat([F.interpolate(cheap_approximation(x).unsqueeze(0), scale_factor=opt_f, mode='nearest-exact') for x in z], dim=0).cpu()
-        #except: pass
-        # Free memory of input latent tensor
-        del z
-
-        # Task queue execution
-        pbar = tqdm(total=num_tiles * len(task_queues[0]), desc=f"[Tiled VAE]: Executing {'Decoder' if is_decoder else 'Encoder'} Task Queue: ")
-
-        # execute the task back and forth when switch tiles so that we always
-        # keep one tile on the GPU to reduce unnecessary data transfer
-        forward = True
-        interrupted = False
-        #state.interrupted = interrupted
-        while True:
-            #if state.interrupted: interrupted = True ; break
-
-            group_norm_param = GroupNormParam()
-            for i in range(num_tiles) if forward else reversed(range(num_tiles)):
-                #if state.interrupted: interrupted = True ; break
-
-                tile = tiles[i].to(device)
-                input_bbox = in_bboxes[i]
-                task_queue = task_queues[i]
-
-                interrupted = False
-                while len(task_queue) > 0:
-                    #if state.interrupted: interrupted = True ; break
-
-                    # DEBUG: current task
-                    # print('Running task: ', task_queue[0][0], ' on tile ', i, '/', num_tiles, ' with shape ', tile.shape)
-                    task = task_queue.pop(0)
-                    if task[0] == 'pre_norm':
-                        group_norm_param.add_tile(tile, task[1])
-                        break
-                    elif task[0] == 'store_res' or task[0] == 'store_res_cpu':
-                        task_id = 0
-                        res = task[1](tile)
-                        if not self.fast_mode or task[0] == 'store_res_cpu':
-                            res = res.cpu()
-                        while task_queue[task_id][0] != 'add_res':
-                            task_id += 1
-                        task_queue[task_id][1] = res
-                    elif task[0] == 'add_res':
-                        tile += task[1].to(device)
-                        task[1] = None
-                    else:
-                        tile = task[1](tile)
-                    pbar.update(1)
-
-                if interrupted: break
-
-                # check for NaNs in the tile.
-                # If there are NaNs, we abort the process to save user's time
-                #devices.test_for_nans(tile, "vae")
-
-                #print(tiles[i].shape, tile.shape, i, num_tiles)
-                if len(task_queue) == 0:
-                    tiles[i] = None
-                    num_completed += 1
-                    if result is None:      # NOTE: dim C varies from different cases, can only be inited dynamically
-                        result = torch.zeros((N, tile.shape[1], height * 8 if is_decoder else height // 8, width * 8 if is_decoder else width // 8), device=device, requires_grad=False)
-                    result[:, :, out_bboxes[i][2]:out_bboxes[i][3], out_bboxes[i][0]:out_bboxes[i][1]] = crop_valid_region(tile, in_bboxes[i], out_bboxes[i], is_decoder)
-                    del tile
-                elif i == num_tiles - 1 and forward:
-                    forward = False
-                    tiles[i] = tile
-                elif i == 0 and not forward:
-                    forward = True
-                    tiles[i] = tile
-                else:
-                    tiles[i] = tile.cpu()
-                    del tile
-
-            if interrupted: break
-            if num_completed == num_tiles: break
-
-            # insert the group norm task to the head of each task queue
-            group_norm_func = group_norm_param.summary()
-            if group_norm_func is not None:
-                for i in range(num_tiles):
-                    task_queue = task_queues[i]
-                    task_queue.insert(0, ('apply_norm', group_norm_func))
-
-        # Done!
-        pbar.close()
-        return result if result is not None else result_approx.to(device)

video.py

@@ -1,57 +0,0 @@
-"""Render the zoom levels into one continuous zoom clip.
-
-Each segment pushes in on level i until its frame is exactly the blurry crop that level i+1
-was made from, then holds still and cross-fades blurry into sharp. The hold is the point:
-it is where the viewer sees what the model actually added.
-"""
-import tempfile
-
-import imageio.v2 as imageio
-import numpy as np
-from PIL import Image
-
-from geometry import window_rect
-
-FPS = 30
-PUSH_FRAMES = 30
-FADE_FRAMES = 14
-TAIL_FRAMES = 24
-
-
-def _ease(t):
-    return t * t * (3 - 2 * t)
-
-
-def _lerp_rect(a, b, t):
-    return tuple(round(p + (q - p) * t) for p, q in zip(a, b))
-
-
-def _frames(levels, upscale, center):
-    size = levels[0].size
-    full = (0, 0, size[0], size[1])
-    for i in range(len(levels) - 1):
-        target = window_rect(size, upscale, center)
-        for f in range(PUSH_FRAMES):
-            # f / (PUSH_FRAMES - 1), not f / PUSH_FRAMES: the last frame has to land exactly on
-            # `target`, which is the crop the next level was built from. Otherwise it jumps.
-            rect = _lerp_rect(full, target, _ease(f / (PUSH_FRAMES - 1)))
-            yield levels[i].crop(rect).resize(size, Image.BICUBIC)
-        blurry = np.asarray(levels[i].crop(target).resize(size, Image.BICUBIC), dtype=np.float32)
-        sharp = np.asarray(levels[i + 1], dtype=np.float32)
-        for f in range(FADE_FRAMES):
-            a = _ease((f + 1) / FADE_FRAMES)
-            yield Image.fromarray((blurry * (1 - a) + sharp * a).astype(np.uint8))
-    for _ in range(TAIL_FRAMES):
-        yield levels[-1]
-
-
-def render(levels, upscale=4, center=(0.5, 0.5)):
-    """Write an mp4 of the zoom and return its path. Needs at least two levels."""
-    if len(levels) < 2:
-        return None
-    path = tempfile.mkstemp(suffix=".mp4")[1]
-    with imageio.get_writer(path, fps=FPS, codec="libx264", quality=8,
-                            pixelformat="yuv420p", macro_block_size=1) as out:
-        for frame in _frames(levels, upscale, center):
-            out.append_data(np.asarray(frame.convert("RGB")))
-    return path

zoom.py

@@ -1,138 +0,0 @@
-"""OracleZoom recursive zoom, one image at a time, all in memory.
-
-Flattened from src/opd_zoom/teacher/oracle_infer.py for the Space. Three differences, none
-of them to the model: no disk round-trip between recursions, the zoom window can sit
-anywhere instead of only the centre, and the loop is a generator so the UI can show each
-level the moment it lands. Everything on one device, which is what ZeroGPU gives us.
-"""
-import gc
-import os
-import sys
-
-import torch
-from PIL import Image
-from torchvision import transforms
-
-from geometry import PROCESS_SIZE, resize_and_center_crop, zoom_window
-
-VENDOR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vendor")
-if VENDOR not in sys.path:
-    sys.path.insert(0, VENDOR)
-
-REPO = "dipta007/OracleZoom"
-SD3 = "stabilityai/stable-diffusion-3-medium-diffusers"
-VLM = "Qwen/Qwen2.5-VL-3B-Instruct"
-COZ_PROMPT = (
-    "The second image is a zoom-in of the first image. Based on this knowledge, "
-    "what is in the second image? Give me a set of words."
-)
-_to_tensor = transforms.Compose([transforms.ToTensor()])
-
-
-class _SRArgs:
-    def __init__(self, coz_ckpt):
-        self.lora_path = f"{coz_ckpt}/SR_LoRA/model_20001.pkl"
-        self.vae_path = f"{coz_ckpt}/SR_VAE/vae_encoder_20001.pt"
-        self.pretrained_model_name_or_path = SD3
-        self.process_size = PROCESS_SIZE
-        self.lora_rank = 4
-        self.merge_and_unload_lora = False
-        self.mixed_precision = "fp16"
-
-
-class Models:
-    """Built once at import, read-only afterwards. Concurrent requests only read."""
-
-    def __init__(self, weights=None):
-        from huggingface_hub import snapshot_download
-
-        w = weights or snapshot_download(repo_id=REPO)
-        self.vlm, self.proc, self.process_vision_info = _build_vlm(f"{w}/ckpt/VLM_LoRA/checkpoint-10000")
-        self.sr = _build_sr(f"{w}/ckpt", f"{w}/merged_transformer.safetensors")
-
-
-def _build_vlm(lora_path):
-    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-    from qwen_vl_utils import process_vision_info
-    from peft import PeftModel
-
-    # device_map="auto" asks accelerate to read real GPU memory, which does not exist at
-    # module scope on ZeroGPU. Place it ourselves. sdpa because flash-attn has no wheel here.
-    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-        VLM, torch_dtype=torch.bfloat16, attn_implementation="sdpa"
-    )
-    # torch_device="cpu" is load-bearing. Left to itself peft calls infer_device(), which reads
-    # the cuda flag ZeroGPU patches to True, then hands "cuda" to safetensors. That builds
-    # tensors straight on a device that is not attached yet: RuntimeError, no CUDA GPUs.
-    model = PeftModel.from_pretrained(model, lora_path, torch_device="cpu").merge_and_unload()
-    return model.eval().to("cuda"), AutoProcessor.from_pretrained(VLM), process_vision_info
-
-
-def _build_sr(coz_ckpt, merged_transformer):
-    from safetensors.torch import load_file
-    from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler
-
-    # Assemble entirely on CPU, then move once. Two places here materialise weights straight
-    # onto the model's own device (inject_lora's torch.load map_location, and safetensors), and
-    # on ZeroGPU that device reads as cuda while no GPU is attached yet.
-    sr = SD3Euler(device="cpu")
-    # Construct first, load second. OSEDiff_SD3_TEST swaps every targeted Linear for a
-    # LoraInjectedLinear, which renames the keys. Our merged checkpoint was saved after that
-    # swap, so loading it earlier matches nothing and silently leaves the base transformer.
-    test = OSEDiff_SD3_TEST(_SRArgs(coz_ckpt), sr)
-    sd = load_file(merged_transformer)
-    missing, unexpected = test.model.transformer.load_state_dict(
-        {k: v.to(torch.float32) for k, v in sd.items()}, strict=False)
-    if len(unexpected) > len(sd) // 2:
-        raise RuntimeError(f"merged transformer did not match: {len(unexpected)} of {len(sd)} "
-                           f"keys unexpected. LoRA injection order or checkpoint is wrong.")
-    print(f"#### merged transformer loaded (missing {len(missing)} unexpected {len(unexpected)})")
-    del sd
-    gc.collect()
-
-    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3):
-        m.to("cuda")
-    sr.transformer.to("cuda", dtype=torch.float32)
-    sr.vae.to("cuda", dtype=torch.float32)
-    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae):
-        m.requires_grad_(False)
-    sr.device = "cuda"   # encode_prompt and set_timesteps read this, not the module devices
-    return test
-
-
-def write_prompt(models, first, second, max_new_tokens=32):
-    messages = [
-        {"role": "system", "content": COZ_PROMPT},
-        {"role": "user", "content": [{"type": "image", "image": first},
-                                     {"type": "image", "image": second}]},
-    ]
-    text = models.proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    images, videos = models.process_vision_info(messages)
-    inputs = models.proc(text=[text], images=images, videos=videos,
-                         padding=True, return_tensors="pt").to("cuda")
-    gen = models.vlm.generate(**inputs, max_new_tokens=max_new_tokens)
-    trimmed = [o[len(i):] for i, o in zip(inputs.input_ids, gen)]
-    return models.proc.batch_decode(trimmed, skip_special_tokens=True,
-                                    clean_up_tokenization_spaces=False)[0].strip()
-
-
-def super_resolve(models, img, prompt):
-    lq = _to_tensor(img).unsqueeze(0).to("cuda") * 2 - 1
-    with torch.no_grad():
-        out = torch.clamp(models.sr(lq, prompt=prompt)[0].cpu(), -1.0, 1.0)
-    return transforms.ToPILImage()(out * 0.5 + 0.5)
-
-
-def zoom(models, image, levels=4, upscale=4, center=(0.5, 0.5)):
-    """Yield one (level, factor, prompt, blurry_input, result) per recursion.
-
-    `blurry_input` is the plain bicubic enlargement the model starts from. It is what the
-    super-resolution has to improve on, so it doubles as the honest before-picture.
-    """
-    cur = resize_and_center_crop(image)
-    yield 0, 1, "", cur, cur
-    for i in range(levels):
-        blurry = zoom_window(cur, upscale, center).resize(cur.size, Image.BICUBIC)
-        prompt = write_prompt(models, cur, blurry)
-        cur = super_resolve(models, blurry, prompt)
-        yield i + 1, upscale ** (i + 1), prompt, blurry, cur