Add the OracleZoom zoom demo

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/0064.png filter=lfs diff=lfs merge=lfs -text
+samples/0245.png filter=lfs diff=lfs merge=lfs -text
+samples/0393.png filter=lfs diff=lfs merge=lfs -text
+samples/0457.png filter=lfs diff=lfs merge=lfs -text
+samples/0479.png filter=lfs diff=lfs merge=lfs -text
+samples/example_16x.png filter=lfs diff=lfs merge=lfs -text
+samples/example_1x.png filter=lfs diff=lfs merge=lfs -text
+samples/example_256x.png filter=lfs diff=lfs merge=lfs -text
+samples/example_256x_input.png filter=lfs diff=lfs merge=lfs -text
+samples/example_4x.png filter=lfs diff=lfs merge=lfs -text
+samples/example_64x.png filter=lfs diff=lfs merge=lfs -text
+samples/example_zoom.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,59 @@
 ---
 title: OracleZoom
-emoji: 🐨
-colorFrom: gray
-colorTo: red
+emoji: 🔎
+colorFrom: yellow
+colorTo: gray
 sdk: gradio
-sdk_version: 6.26.0
-python_version: '3.12'
+sdk_version: 5.50.0
+python_version: "3.12.12"
 app_file: app.py
-pinned: false
+startup_duration_timeout: 45m
+pinned: true
+license: mit
+short_description: Zoom any photo to 256x, one 4x step at a time
+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
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# OracleZoom
+
+Upload a photo, pick a point, and watch it zoom to 256x. Four steps of 4x, each one drawn
+from the last.
+
+- Paper: https://arxiv.org/abs/2609.06490
+- Code: https://github.com/dipta007/OracleZoom
+- Project page: https://dipta007.github.io/OracleZoom/
+- Model: https://huggingface.co/dipta007/OracleZoom
+
+## Running it yourself
+
+The Space needs an `HF_TOKEN` secret. Stable Diffusion 3-medium is gated, so the token has
+to come from an account that has accepted its licence. Nothing in the code reads the token
+directly; `huggingface_hub` picks up `HF_TOKEN` on its own.
+
+Hardware: ZeroGPU. The pipeline holds about 27 GB of weights, well inside the 48 GB that a
+`large` slice gives. ZeroGPU is Gradio-only, so a Docker Space cannot host this.
+
+## What is in here
+
+`app.py` is the interface. `zoom.py` is the recursion, flattened from the research repo so it
+runs in memory with no disk round-trips and lets the zoom window sit off centre. `video.py`
+renders the clip. `vendor/` holds four files copied from
+[Chain-of-Zoom](https://github.com/bryanswkim/Chain-of-Zoom) (MIT, see `vendor/COZ-LICENSE`),
+because Spaces do not support git submodules. The only edit to them pins `weights_only=True`
+on three `torch.load` calls, whose default flipped in torch 2.6.
+
+## 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 says which claims we do and do not make.

app.py

@@ -0,0 +1,187 @@
+"""OracleZoom demo: upload a photo, watch it zoom 4x to 256x."""
+import gradio as gr
+import spaces
+from PIL import Image, ImageDraw
+
+import geometry
+import video
+import zoom
+
+MODELS = zoom.Models()
+
+UPSCALE = 4
+ACCENT = "#f5b942"
+LABELS = {1: "input", 4: "4x", 16: "16x", 64: "64x", 256: "256x"}
+SAMPLES = ["0479", "0064", "0245", "0393", "0457"]
+
+# Shipped output of a real run, so the page shows the payoff before anyone spends any quota.
+EXAMPLE_CLIP = "samples/example_zoom.mp4"
+EXAMPLE_LEVELS = [("samples/example_1x.png", "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 = "_An example run. Upload a photo above to make your own._"
+
+HEADER = """
+<div style="text-align:center;max-width:820px;margin:0 auto 4px">
+  <h1 style="margin:0;font-size:2.1em;letter-spacing:-.02em">OracleZoom 🔎</h1>
+  <p style="margin:.5em 0 .9em;font-size:1.08em;line-height:1.5;opacity:.85">
+    Zoom into any photo far past what it holds. Four steps of 4x take you to
+    <b>256x</b>, each one drawn from the last.
+  </p>
+  <p style="margin:0;font-size:.95em">
+    <a href="https://arxiv.org/abs/2609.06490">Paper</a> &nbsp;·&nbsp;
+    <a href="https://github.com/dipta007/OracleZoom">Code</a> &nbsp;·&nbsp;
+    <a href="https://dipta007.github.io/OracleZoom/">Project page</a> &nbsp;·&nbsp;
+    <a href="https://huggingface.co/dipta007/OracleZoom">Model</a>
+  </p>
+</div>
+"""
+
+HOW = """
+##### How it goes
+
+1. Crop your photo to a 512 square.
+2. Zoom 4x into the amber box. That crop is blurry, so a vision language model describes it.
+3. Super-resolve it with that description as the guide.
+4. Repeat on the result. Four rounds reach 256x.
+
+Runs free on ZeroGPU, so the first zoom after a quiet spell waits for a GPU.
+"""
+
+FOOTER = """
+### What you are looking at
+
+Standard super-resolution models break down well before 16x. OracleZoom gets to 256x by
+zooming one 4x step at a time and feeding each result into the next step. At every step a
+vision language model writes a short description of the crop, and that description guides
+the detail the super-resolution model draws.
+
+Past the first step or two there is no ground truth to recover, so the deep levels are
+**plausible detail, not measured detail**. The zoom point is a crop of your photo, not a
+real camera lens moving closer. Read the paper for what we do and do not claim.
+
+```bibtex
+@inproceedings{dipta2027oraclezoom,
+  title  = {OracleZoom: Reference-Constrained Recursive Super-Resolution},
+  author = {Roy Dipta, Shubhashis and Saha, Sourajit and Saha, Shaswati and Sarwar, Nobin},
+  year   = {2027}
+}
+```
+"""
+
+CSS = """
+#hero video {border-radius:12px}
+.contain {max-width:1400px !important}
+footer {display:none !important}
+"""
+
+
+def preview(image, levels, cx, cy):
+    """Show where the zoom will go, so nobody spends a run to find out."""
+    if image is None:
+        return None
+    canvas = geometry.resize_and_center_crop(image).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)
+    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):
+    x, y = evt.index
+    return x / geometry.PROCESS_SIZE, y / geometry.PROCESS_SIZE
+
+
+def estimate_duration(image, levels, cx, cy):
+    # 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 mid-way. Provisional until
+    # measured on the real hardware.
+    return int(15 + 9 * int(levels))
+
+
+def _stream(image, levels, cx, cy):
+    if image is None:
+        raise gr.Error("Upload a photo first.")
+    levels = int(levels)
+    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, LABELS.get(factor, f"{factor}x")))
+        if step:
+            compare = (blurry, result)
+            prompts.append(f"**{LABELS[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 **{LABELS[factor]}**, {left} step{'s' if left > 1 else ''} to go…"
+        else:
+            status = f"At **{LABELS[factor]}**. Rendering the clip…"
+        # 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)
+    yield ("Done. 🔎", video.render(frames, UPSCALE, (cx, cy)), gallery, compare,
+           "\n\n".join(prompts))
+
+
+@spaces.GPU(duration=estimate_duration)
+def run(image, levels, cx, cy):
+    yield from _stream(image, levels, cx, cy)
+
+
+@spaces.GPU(duration=estimate_duration(None, 4, 0, 0))
+def run_example(image):
+    """One input, so the examples render as thumbnails instead of a four-column table.
+    Its results are cached, which is how a visitor with no quota left still sees output."""
+    yield from _stream(image, 4, 0.5, 0.5)
+
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber"), css=CSS,
+               title="OracleZoom: zoom to 256x") as demo:
+    gr.HTML(HEADER)
+    with gr.Row():
+        with gr.Column(scale=4):
+            image = gr.Image(label="Your photo", type="pil", height=300, sources=["upload", "clipboard"])
+            # Seeded with the example so the box overlay explains itself before any upload.
+            target = gr.Image(value=preview(Image.open(EXAMPLE_LEVELS[0][0]), 4, 0.5, 0.5),
+                              label="Click to move the zoom point", type="pil",
+                              interactive=False, height=340, show_download_button=False)
+            levels = gr.Slider(1, 4, value=4, step=1, label="Zoom steps",
+                               info="1 step = 4x, 4 steps = 256x. Fewer steps finish sooner.")
+            with gr.Accordion("Set the point by hand", open=False):
+                cx = gr.Slider(0, 1, value=0.5, step=0.01, label="Horizontal")
+                cy = gr.Slider(0, 1, value=0.5, step=0.01, label="Vertical")
+            go = gr.Button("🔎 Zoom in", variant="primary", size="lg")
+            gr.Markdown(HOW)
+        with gr.Column(scale=6):
+            status = gr.Markdown("Upload a photo, then press **Zoom in**.")
+            clip = gr.Video(value=EXAMPLE_CLIP, label="The zoom", elem_id="hero", autoplay=True,
+                            loop=True, show_share_button=True, height=400)
+            gallery = gr.Gallery(value=EXAMPLE_LEVELS, label="Every level", columns=5, height=175,
+                                 object_fit="cover", show_download_button=True)
+            compare = gr.ImageSlider(value=EXAMPLE_COMPARE, height=400,
+                                     label="256x: plain enlargement (left) vs OracleZoom (right)")
+            gr.Markdown("##### What the model said it saw, step by step")
+            prompts = gr.Markdown(EXAMPLE_NOTE)
+
+    controls = [image, levels, cx, cy]
+    for c in controls:
+        c.change(preview, controls, target, show_api=False)
+    target.select(pick_point, None, [cx, cy], show_api=False)
+    go.click(run, controls, [status, clip, gallery, compare, prompts])
+
+    gr.Examples(
+        examples=[f"samples/{n}.png" for n in SAMPLES],
+        inputs=[image],
+        outputs=[status, clip, gallery, compare, prompts],
+        fn=run_example,
+        cache_examples=True,
+        label="Or try one of these (already computed, costs you nothing)",
+    )
+    gr.Markdown(FOOTER)
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

geometry.py

@@ -0,0 +1,38 @@
+"""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."""
+    rects, rect = [], (0, 0, size[0], size[1])
+    for _ in range(levels):
+        side = rect[2] - rect[0]
+        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

requirements.txt

@@ -0,0 +1,21 @@
+# ZeroGPU only runs torch 2.8.0 and newer, which is why these pins differ from the repo's
+# pyproject.toml. 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.
+# `spaces` is deliberately absent. The platform pins its own copy and a second pin breaks pip.
+torch==2.8.0
+torchvision==0.23.0
+diffusers==0.32.2
+transformers==4.49.0
+accelerate
+peft>=0.13
+safetensors
+huggingface-hub
+qwen-vl-utils
+lpips
+einops
+pyyaml
+tqdm
+numpy
+pillow
+imageio
+imageio-ffmpeg

samples/example_zoom.mp4

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

test_app_smoke.py

@@ -0,0 +1,91 @@
+"""Build the whole Gradio graph with the models stubbed out.
+
+Catches component, parameter and event-wiring mistakes without a GPU, which is most of what
+breaks a Space on first deploy. The model path itself is not covered here.
+
+    uv run --no-project --python 3.12 --with pytest --with "gradio==5.50.0" --with spaces \
+        --with pillow --with numpy --with imageio --with imageio-ffmpeg \
+        pytest space/test_app_smoke.py -q
+"""
+import os
+import sys
+import types
+
+import numpy as np
+import pytest
+from PIL import Image
+
+
+@pytest.fixture(scope="module")
+def app():
+    """Import app.py with `zoom` replaced by a stub that skips the GPU work."""
+    import geometry
+
+    stub = types.ModuleType("zoom")
+    stub.REPO = "dipta007/OracleZoom"
+    stub.Models = type("Models", (), {"__init__": lambda self, weights=None: None})
+
+    def fake_zoom(models, image, levels=4, upscale=4, center=(0.5, 0.5)):
+        cur = geometry.resize_and_center_crop(image)
+        yield 0, 1, "", cur, cur
+        for i in range(levels):
+            blurry = geometry.zoom_window(cur, upscale, center).resize(cur.size, Image.BICUBIC)
+            yield i + 1, upscale ** (i + 1), f"tag {i}", blurry, blurry
+            cur = blurry
+
+    stub.zoom = fake_zoom
+    sys.modules["zoom"] = stub
+    import app as module
+    return module
+
+
+@pytest.fixture
+def photo():
+    rng = np.random.default_rng(0)
+    return Image.fromarray(rng.integers(0, 255, (700, 900, 3), dtype=np.uint8))
+
+
+def test_the_interface_builds(app):
+    assert app.demo.blocks and app.demo.fns
+
+
+def test_preview_returns_a_canvas_and_tolerates_no_image(app, photo):
+    assert app.preview(photo, 4, 0.3, 0.7).size == (512, 512)
+    assert app.preview(None, 4, 0.5, 0.5) is None
+
+
+def test_declared_duration_grows_with_depth(app, photo):
+    got = [app.estimate_duration(photo, n, 0.5, 0.5) for n in (1, 2, 3, 4)]
+    assert got == sorted(got) and all(isinstance(v, int) for v in got)
+
+
+def test_the_shipped_example_assets_exist(app):
+    for path in [app.EXAMPLE_CLIP, *app.EXAMPLE_COMPARE, *(p for p, _ in app.EXAMPLE_LEVELS)]:
+        assert os.path.getsize(path) > 1000, path
+    assert [label for _, label in app.EXAMPLE_LEVELS] == ["input", "4x", "16x", "64x", "256x"]
+
+
+def test_every_sample_named_for_the_examples_is_present(app):
+    for name in app.SAMPLES:
+        assert os.path.exists(f"samples/{name}.png")
+
+
+def test_run_streams_one_update_per_level_then_the_clip(app, photo):
+    yields = list(app.run(photo, 3, 0.5, 0.5))
+    assert len(yields) == 5                       # entry level, 3 zoom steps, then the clip
+    assert all(len(y) == 5 for y in yields)
+    assert yields[0][1] is None                   # first yield clears any stale clip
+    status, clip, gallery, compare, prompts = yields[-1]
+    assert [label for _, label in gallery] == ["input", "4x", "16x", "64x"]
+    assert compare is not None and len(compare) == 2
+    assert prompts.count("tag") == 3
+    try:
+        assert os.path.getsize(clip) > 10_000
+    finally:
+        os.remove(clip)
+
+
+def test_run_rejects_an_empty_upload(app):
+    import gradio as gr
+    with pytest.raises(gr.Error):
+        next(app.run(None, 4, 0.5, 0.5))

test_space.py

@@ -0,0 +1,90 @@
+"""Tests for the parts of the Space that need no GPU: crop geometry and the video render.
+
+    uv run --with pytest --with pillow --with numpy --with imageio --with imageio-ffmpeg \
+        pytest space/test_space.py -q
+"""
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+
+import geometry
+import video
+
+
+def _img(seed, size=512):
+    rng = np.random.default_rng(seed)
+    return Image.fromarray(rng.integers(0, 255, (size, size, 3), dtype=np.uint8))
+
+
+def test_entry_crop_is_square_and_centred():
+    out = geometry.resize_and_center_crop(_img(0, 900).resize((1200, 900)))
+    assert out.size == (512, 512)
+
+
+def test_window_is_one_quarter_wide_and_centred_by_default():
+    assert geometry.window_rect((512, 512), 4) == (192, 192, 320, 320)
+
+
+def test_window_stays_inside_the_frame_at_the_edges():
+    for center in [(0.0, 0.0), (1.0, 1.0), (0.0, 1.0), (1.0, 0.0)]:
+        l, t, r, b = geometry.window_rect((512, 512), 4, center)
+        assert 0 <= l < r <= 512 and 0 <= t < b <= 512
+        assert (r - l, b - t) == (128, 128)
+
+
+def test_nested_rects_shrink_by_upscale_and_nest():
+    rects = geometry.nested_rects((512, 512), 4, 4)
+    assert [r[2] - r[0] for r in rects] == [128, 32, 8, 2]
+    for inner, outer in zip(rects[1:], rects):
+        assert outer[0] <= inner[0] and inner[2] <= outer[2]
+        assert outer[1] <= inner[1] and inner[3] <= outer[3]
+
+
+def test_zoom_window_matches_window_rect():
+    img = _img(1)
+    assert geometry.zoom_window(img, 4).size == (128, 128)
+    assert np.array_equal(np.asarray(geometry.zoom_window(img, 4)),
+                          np.asarray(img.crop(geometry.window_rect(img.size, 4))))
+
+
+def test_video_needs_two_levels():
+    assert video.render([_img(2)]) is None
+
+
+def test_video_renders_an_mp4_with_the_expected_frame_count():
+    levels = [_img(i) for i in range(5)]
+    path = video.render(levels, 4, (0.5, 0.5))
+    try:
+        assert os.path.getsize(path) > 10_000
+        import imageio.v2 as imageio
+        frames = imageio.mimread(path, memtest=False)
+        expected = 4 * (video.PUSH_FRAMES + video.FADE_FRAMES) + video.TAIL_FRAMES
+        assert len(frames) == expected
+        assert frames[0].shape[:2] == (512, 512)
+    finally:
+        os.remove(path)
+
+
+def test_the_push_ends_exactly_on_the_next_level_input():
+    """The last push frame must equal the blurry crop the next level was built from.
+    If it does not, the clip jumps at every transition."""
+    levels = [_img(3), _img(4)]
+    frames = list(video._frames(levels, 4, (0.5, 0.5)))
+    last_push = np.asarray(frames[video.PUSH_FRAMES - 1], dtype=np.int16)
+    target = geometry.window_rect(levels[0].size, 4, (0.5, 0.5))
+    blurry = np.asarray(levels[0].crop(target).resize(levels[0].size, Image.BICUBIC),
+                        dtype=np.int16)
+    assert np.array_equal(last_push, blurry)
+
+
+def test_the_fade_ends_on_the_sharp_level():
+    levels = [_img(5), _img(6)]
+    frames = list(video._frames(levels, 4, (0.5, 0.5)))
+    end_of_fade = np.asarray(frames[video.PUSH_FRAMES + video.FADE_FRAMES - 1], dtype=np.int16)
+    assert np.abs(end_of_fade - np.asarray(levels[1], dtype=np.int16)).max() <= 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-q"]))

vendor/COZ-LICENSE

@@ -0,0 +1,21 @@
+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

@@ -0,0 +1,137 @@
+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

@@ -0,0 +1,61 @@
+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

@@ -0,0 +1,913 @@
+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

@@ -0,0 +1,138 @@
+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

@@ -0,0 +1,829 @@
+# ------------------------------------------------------------------------
+#
+#   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

@@ -0,0 +1,57 @@
+"""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

@@ -0,0 +1,127 @@
+"""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 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"
+    )
+    model = PeftModel.from_pretrained(model, lora_path).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
+
+    sr = SD3Euler()
+    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)
+    # 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)})")
+    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