Initial commit

.gitattributes

@@ -33,3 +33,4 @@ 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
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,18 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Output
+output*/
+out*/
+out/
+.gradio/
+/*.png
+/*.csv

LICENSE.md

@@ -0,0 +1,25 @@
+STABILITY AI COMMUNITY LICENSE AGREEMENT Last Updated: July 5, 2024
+
+I. INTRODUCTION
+
+This Agreement applies to any individual person or entity ("You", "Your" or "Licensee") that uses or distributes any portion or element of the Stability AI Materials or Derivative Works thereof for any Research & Non-Commercial or Commercial purpose. Capitalized terms not otherwise defined herein are defined in Section V below.
+
+This Agreement is intended to allow research, non-commercial, and limited commercial uses of the Models free of charge. In order to ensure that certain limited commercial uses of the Models continue to be allowed, this Agreement preserves free access to the Models for people or organizations generating annual revenue of less than US $1,000,000 (or local currency equivalent).
+
+By clicking "I Accept" or by using or distributing or using any portion or element of the Stability Materials or Derivative Works, You agree that You have read, understood and are bound by the terms of this Agreement. If You are acting on behalf of a company, organization or other entity, then "You" includes you and that entity, and You agree that You: (i) are an authorized representative of such entity with the authority to bind such entity to this Agreement, and (ii) You agree to the terms of this Agreement on that entity's behalf.
+
+II. RESEARCH & NON-COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement, Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI's intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Research or Non-Commercial Purpose. "Research Purpose" means academic or scientific advancement, and in each case, is not primarily intended for commercial advantage or monetary compensation to You or others. "Non-Commercial Purpose" means any purpose other than a Research Purpose that is not primarily intended for commercial advantage or monetary compensation to You or others, such as personal use (i.e., hobbyist) or evaluation and testing.
+
+III. COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement (including the remainder of this Section III), Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI's intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Commercial Purpose. "Commercial Purpose" means any purpose other than a Research Purpose or Non-Commercial Purpose that is primarily intended for commercial advantage or monetary compensation to You or others, including but not limited to, (i) creating, modifying, or distributing Your product or service, including via a hosted service or application programming interface, and (ii) for Your business's or organization's internal operations. If You are using or distributing the Stability AI Materials for a Commercial Purpose, You must register with Stability AI at (https://stability.ai/community-license). If at any time You or Your Affiliate(s), either individually or in aggregate, generate more than USD $1,000,000 in annual revenue (or the equivalent thereof in Your local currency), regardless of whether that revenue is generated directly or indirectly from the Stability AI Materials or Derivative Works, any licenses granted to You under this Agreement shall terminate as of such date. You must request a license from Stability AI at (https://stability.ai/enterprise) , which Stability AI may grant to You in its sole discretion. If you receive Stability AI Materials, or any Derivative Works thereof, from a Licensee as part of an integrated end user product, then Section III of this Agreement will not apply to you.
+
+IV. GENERAL TERMS
+
+Your Research, Non-Commercial, and Commercial License(s) under this Agreement are subject to the following terms. a. Distribution & Attribution. If You distribute or make available the Stability AI Materials or a Derivative Work to a third party, or a product or service that uses any portion of them, You shall: (i) provide a copy of this Agreement to that third party, (ii) retain the following attribution notice within a "Notice" text file distributed as a part of such copies: "This Stability AI Model is licensed under the Stability AI Community License, Copyright © Stability AI Ltd. All Rights Reserved", and (iii) prominently display "Powered by Stability AI" on a related website, user interface, blogpost, about page, or product documentation. If You create a Derivative Work, You may add your own attribution notice(s) to the "Notice" text file included with that Derivative Work, provided that You clearly indicate which attributions apply to the Stability AI Materials and state in the "Notice" text file that You changed the Stability AI Materials and how it was modified. b. Use Restrictions. Your use of the Stability AI Materials and Derivative Works, including any output or results of the Stability AI Materials or Derivative Works, must comply with applicable laws and regulations (including Trade Control Laws and equivalent regulations) and adhere to the Documentation and Stability AI's AUP, which is hereby incorporated by reference. Furthermore, You will not use the Stability AI Materials or Derivative Works, or any output or results of the Stability AI Materials or Derivative Works, to create or improve any foundational generative AI model (excluding the Models or Derivative Works). c. Intellectual Property. (i) Trademark License. No trademark licenses are granted under this Agreement, and in connection with the Stability AI Materials or Derivative Works, You may not use any name or mark owned by or associated with Stability AI or any of its Affiliates, except as required under Section IV(a) herein. (ii) Ownership of Derivative Works. As between You and Stability AI, You are the owner of Derivative Works You create, subject to Stability AI's ownership of the Stability AI Materials and any Derivative Works made by or for Stability AI. (iii) Ownership of Outputs. As between You and Stability AI, You own any outputs generated from the Models or Derivative Works to the extent permitted by applicable law. (iv) Disputes. If You or Your Affiliate(s) institute litigation or other proceedings against Stability AI (including a cross-claim or counterclaim in a lawsuit) alleging that the Stability AI Materials, Derivative Works or associated outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by You, then any licenses granted to You under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Stability AI from and against any claim by any third party arising out of or related to Your use or distribution of the Stability AI Materials or Derivative Works in violation of this Agreement. (v) Feedback. From time to time, You may provide Stability AI with verbal and/or written suggestions, comments or other feedback related to Stability AI's existing or prospective technology, products or services (collectively, "Feedback"). You are not obligated to provide Stability AI with Feedback, but to the extent that You do, You hereby grant Stability AI a perpetual, irrevocable, royalty-free, fully-paid, sub-licensable, transferable, non-exclusive, worldwide right and license to exploit the Feedback in any manner without restriction. Your Feedback is provided "AS IS" and You make no warranties whatsoever about any Feedback. d. Disclaimer Of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE STABILITY AI MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OR LAWFULNESS OF USING OR REDISTRIBUTING THE STABILITY AI MATERIALS, DERIVATIVE WORKS OR ANY OUTPUT OR RESULTS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE STABILITY AI MATERIALS, DERIVATIVE WORKS AND ANY OUTPUT AND RESULTS. e. Limitation Of Liability. IN NO EVENT WILL STABILITY AI OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT, INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF STABILITY AI OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING. f. Term And Termination. The term of this Agreement will commence upon Your acceptance of this Agreement or access to the Stability AI Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Stability AI may terminate this Agreement if You are in breach of any term or condition of this Agreement. Upon termination of this Agreement, You shall delete and cease use of any Stability AI Materials or Derivative Works. Section IV(d), (e), and (g) shall survive the termination of this Agreement. g. Governing Law. This Agreement will be governed by and constructed in accordance with the laws of the United States and the State of California without regard to choice of law principles, and the UN Convention on Contracts for International Sale of Goods does not apply to this Agreement.
+
+V. DEFINITIONS
+
+"Affiliate(s)" means any entity that directly or indirectly controls, is controlled by, or is under common control with the subject entity; for purposes of this definition, "control" means direct or indirect ownership or control of more than 50% of the voting interests of the subject entity. "Agreement" means this Stability AI Community License Agreement. "AUP" means the Stability AI Acceptable Use Policy available at https://stability.ai/use-policy, as may be updated from time to time. "Derivative Work(s)" means (a) any derivative work of the Stability AI Materials as recognized by U.S. copyright laws and (b) any modifications to a Model, and any other model created which is based on or derived from the Model or the Model's output, including"fine tune" and "low-rank adaptation" models derived from a Model or a Model's output, but do not include the output of any Model. "Documentation" means any specifications, manuals, documentation, and other written information provided by Stability AI related to the Software or Models. "Model(s)" means, collectively, Stability AI's proprietary models and algorithms, including machine-learning models, trained model weights and other elements of the foregoing listed on Stability's Core Models Webpage available at, https://stability.ai/core-models, as may be updated from time to time. "Stability AI" or "we" means Stability AI Ltd. and its Affiliates. "Software" means Stability AI's proprietary software made available under this Agreement now or in the future. "Stability AI Materials" means, collectively, Stability's proprietary Models, Software and Documentation (and any portion or combination thereof) made available under this Agreement. "Trade Control Laws" means any applicable U.S. and non-U.S. export control and trade sanctions laws and regulations.

README.md

@@ -1,12 +1,32 @@
 ---
 title: ReSWD
-emoji: 🌍
+emoji: 📊
 colorFrom: purple
 colorTo: green
 sdk: gradio
 sdk_version: 5.47.1
 app_file: app.py
 pinned: false
+models:
+  - stabilityai/stable-diffusion-3.5-large-turbo
+license: other
+license_name: stabilityai-ai-community
+license_link: LICENSE.md
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ReSWD: ReSTIR‘d, not shaken. Combining Reservoir Sampling and Sliced Wasserstein Distance for Variance Reduction.
+
+<a href="https://reservoirswd.github.io/"><img src="https://img.shields.io/badge/Project%20Page-5CE1BC.svg"></a> <a href="https://reservoirswd.github.io/static/paper.pdf"><img src="https://img.shields.io/badge/Arxiv-2408.00653-B31B1B.svg"></a> <a href="https://huggingface.co/spaces/stabilityai/reswd"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Gradio%20Demo-Huggingface-orange"></a>
+
+This is the official codebase for **ReSWD**, a state-of-the-art algorithm for distribution matching with reduced variance. It has several applications (such as diffusion guidance or color matching).
+
+## Citation
+
+```BibTeX
+@article{boss2025reswd,
+  title={ReSWD: ReSTIR‘d, not shaken. Combining Reservoir Sampling and Sliced Wasserstein Distance for Variance Reduction.},
+  author={Boss, Mark and Engelhardt, Andreas and Donné, Simon and Jampani, Varun},
+  journal={arXiv preprint},
+  year={2025}
+}
+```

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+import lightning as L
+
+from src.gradio_demo.color_matching import create_color_matching
+from src.gradio_demo.sw_guidance import create_sw_guidance
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # ReSWD
+
+        ReSTIR‘d, not shaken. Combining Reservoir Sampling and Sliced Wasserstein
+        Distance for Variance Reduction.
+
+        ReSWD is a method for distribution matching with reduced variance.
+        """
+    )
+    fabric = L.Fabric(devices=1, accelerator="auto", precision="16-mixed")
+
+    with gr.Tab("SW Guidance (SD 3.5 Large Turbo)"):
+        create_sw_guidance(fabric, "stabilityai/stable-diffusion-3.5-large-turbo")
+    with gr.Tab("Color Matching"):
+        create_color_matching(fabric)
+
+    demo.launch()

requirements.txt

@@ -0,0 +1,17 @@
+accelerate>=1.7.1
+diffusers>=0.33.1
+imageio>=2.37.0
+jaxtyping>=0.3.2
+lightning>=2.5.1.post0
+lxml>=5.4.0
+matplotlib>=3.10.3
+numpy>=2.2.5
+opencv-python>=4.11.0.86
+pillow>=11.2.1
+protobuf>=6.32.1
+scipy>=1.15.3
+sentencepiece>=0.2.1
+torch>=2.7.0
+torchvision>=0.22.0
+tqdm>=4.67.1
+transformers>=4.52.4

src/color_matcher.py

@@ -0,0 +1,205 @@
+import gc
+import os
+from typing import List, Literal, Optional, Tuple
+
+import imageio
+import lightning as L
+import torch
+from jaxtyping import Float
+from torchvision.transforms import Resize
+from tqdm import tqdm
+
+from src.loss import AbstractLoss
+from src.loss.vector_swd import VectorSWDLoss
+from src.utils.asc_cdl import asc_cdl_forward, save_asc_cdl
+from src.utils.color_space import rgb_to_lab
+from src.utils.image import from_torch, read_img, to_torch, write_img
+
+
+class CDL(torch.nn.Module):
+    def __init__(self, batch_size: int):
+        super().__init__()
+        self.cdl_slope = torch.nn.Parameter(torch.ones(batch_size, 3))
+        self.cdl_offset = torch.nn.Parameter(torch.zeros(batch_size, 3))
+        self.cdl_power = torch.nn.Parameter(torch.ones(batch_size, 3))
+        self.cdl_saturation = torch.nn.Parameter(torch.ones(batch_size))
+
+    def forward(
+        self, x: Float[torch.Tensor, "*B C H W"]
+    ) -> Float[torch.Tensor, "*B C H W"]:
+        return asc_cdl_forward(
+            x, self.cdl_slope, self.cdl_offset, self.cdl_power, self.cdl_saturation
+        )
+
+    def to_cdl_xml(self) -> str:
+        ret = []
+        for b in range(self.cdl_slope.shape[0]):
+            ret.append(
+                save_asc_cdl(
+                    {
+                        "slope": self.cdl_slope[b],
+                        "offset": self.cdl_offset[b],
+                        "power": self.cdl_power[b],
+                        "saturation": self.cdl_saturation[b],
+                    },
+                    None,
+                )
+            )
+        return ret
+
+    def save(self, path: str):
+        for b in range(self.cdl_slope.shape[0]):
+            save_asc_cdl(
+                {
+                    "slope": self.cdl_slope[b],
+                    "offset": self.cdl_offset[b],
+                    "power": self.cdl_power[b],
+                    "saturation": self.cdl_saturation[b],
+                },
+                os.path.join(path, f"cdl_{b}.xml"),
+            )
+
+
+def train(
+    fabric: L.Fabric,
+    criteria: AbstractLoss,
+    source_img: Float[torch.Tensor, "B C H W"],
+    target_img: Float[torch.Tensor, "B C H W"],
+    num_steps: int,
+    lr: float,
+    match_resolution: int,
+    silent: bool = False,
+    write_video_animation_path: Optional[str] = None,
+) -> Tuple[Float[torch.Tensor, "*B C H W"], CDL, List[float]]:
+    criteria = fabric.setup(criteria)
+
+    source_max_res = Resize(match_resolution, antialias=True)(source_img)
+    target_max_res = Resize(match_resolution, antialias=True)(target_img)
+
+    target_cielab = (
+        fabric.to_device(rgb_to_lab(target_max_res).permute(0, 3, 1, 2))
+        .permute(0, 2, 3, 1)
+        .contiguous()
+    )
+
+    source_max_res = fabric.to_device(source_max_res)
+    source_img = fabric.to_device(source_img)
+
+    batch_size = source_img.shape[0]
+    cdl = CDL(batch_size)
+
+    optim = torch.optim.Adam(cdl.parameters(), lr=lr)
+    cdl, optim = fabric.setup(cdl, optim)
+
+    lossses = []
+    for i in tqdm(range(num_steps), disable=silent):
+        optim.zero_grad(set_to_none=True)
+
+        cdl_source = cdl(source_max_res)
+        source_cielab = (
+            rgb_to_lab(cdl_source.permute(0, 3, 1, 2)).permute(0, 2, 3, 1).contiguous()
+        )
+
+        loss = criteria(
+            source_cielab.view(source_cielab.shape[0], source_cielab.shape[1], -1),
+            target_cielab.view(target_cielab.shape[0], target_cielab.shape[1], -1),
+            i,
+        )
+
+        fabric.backward(loss)
+        optim.step()
+
+        lossses.append(loss.item())
+
+        if write_video_animation_path is not None:
+            write_img(
+                os.path.join(write_video_animation_path, f"{i:05d}.jpg"),
+                from_torch(cdl(source_img).squeeze(0) * 2 - 1),
+            )
+
+    source_full_res_cdl = cdl(source_img)
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return source_full_res_cdl, cdl, lossses
+
+
+def run(
+    save_dir: str,
+    source_img: List[str],
+    target_img: List[str],
+    matching_resolution: int,
+    precision: Literal["32-true", "16-mixed"] = "16-mixed",
+    num_projections: int = 64,
+    lr: float = 0.01,
+    steps: int = 300,
+    use_ucv: bool = False,
+    use_lcv: bool = False,
+    distance: Literal["l1", "l2"] = "l1",
+    refresh_projections_every_n_steps: int = 1,
+    num_new_candidates: int = 32,
+    sampling_mode: Literal["gaussian", "qmc"] = "gaussian",
+    write_video: bool = False,
+    **kwargs,
+):
+    fabric = L.Fabric(devices=1, accelerator="auto", precision=precision)
+
+    source_imgs = torch.stack(
+        [to_torch(read_img(s)) * 0.5 + 0.5 for s in source_img], dim=0
+    )
+    target_imgs = torch.stack(
+        [to_torch(read_img(t)) * 0.5 + 0.5 for t in target_img], dim=0
+    )
+
+    criteria = VectorSWDLoss(
+        num_proj=num_projections,
+        distance=distance,
+        use_ucv=use_ucv,
+        use_lcv=use_lcv,
+        refresh_projections_every_n_steps=refresh_projections_every_n_steps,
+        num_new_candidates=num_new_candidates,
+        sampling_mode=sampling_mode,
+    )
+
+    os.makedirs(save_dir, exist_ok=True)
+    animation_dir = os.path.join(save_dir, "animation")
+
+    if write_video:
+        os.makedirs(animation_dir, exist_ok=True)
+
+    source_full_res_cdl, cdl, lossses = train(
+        fabric,
+        criteria,
+        source_imgs,
+        target_imgs,
+        steps,
+        lr,
+        matching_resolution,
+        write_video_animation_path=animation_dir if write_video else None,
+    )
+
+    cdl.save(save_dir)
+
+    for i, img in enumerate(source_full_res_cdl):
+        write_img(
+            os.path.join(save_dir, f"color_matched_{i}.png"),
+            from_torch(img * 2 - 1),
+        )
+
+    if write_video:
+        # Get the list of image files in the animation directory
+        image_files = [f for f in os.listdir(animation_dir) if f.endswith(".jpg")]
+        image_files.sort(
+            key=lambda x: int(x.split(".")[0])
+        )  # Ensure they are in the correct order
+
+        # Create a video from the images
+        with imageio.get_writer(
+            os.path.join(save_dir, "animation.mp4"), fps=30, codec="libx264"
+        ) as writer:
+            for image_file in image_files:
+                image = imageio.imread(os.path.join(animation_dir, image_file))
+                writer.append_data(image)
+
+    return source_full_res_cdl, cdl, lossses

src/gradio_demo/color_matching.py

@@ -0,0 +1,164 @@
+import os
+
+import gradio as gr
+import lightning as L
+import numpy as np
+import spaces
+
+from src.color_matcher import train
+from src.loss import VectorSWDLoss
+from src.utils.image import from_torch, to_torch
+
+
+def create_color_matching(fabric: L.Fabric):
+    """
+    Creates the Gradio interface for color matching between source and target images.
+    """
+    gr.Markdown(
+        """
+        # Color Matching
+        Matches the color of source images to target images using ASC CDL parameters.
+        """
+    )
+
+    gr.Markdown("## Source Images")
+    with gr.Row(variant="compact"):
+        source_image = gr.Image(label="Source Image", height=512)
+        target_image = gr.Image(label="Target Image", height=512)
+    with gr.Row(variant="compact"):
+        example_paths = os.path.join("example", "color_matching")
+        # Find all images in the example directory
+        example_images = [
+            os.path.join(example_paths, f)
+            for f in os.listdir(example_paths)
+            if os.path.isfile(os.path.join(example_paths, f))
+            and os.path.splitext(f)[-1] in [".png", ".jpg", ".jpeg"]
+        ]
+
+        gr.Examples(
+            examples=example_images,
+            inputs=[source_image],
+        )
+        gr.Examples(
+            examples=example_images,
+            inputs=[target_image],
+        )
+
+    gr.Markdown(
+        """
+        # Configuration
+        Adjust the parameters for the color matching process.
+        """
+    )
+
+    with gr.Accordion("Advanced Config", open=False):
+        learning_rate_slider = gr.Slider(
+            1e-4, 0.1, value=1e-3, step=1e-4, label="Learning rate"
+        )
+        match_resolution_slider = gr.Slider(
+            64, 1024, value=128, step=64, label="Match resolution"
+        )
+        num_steps_slider = gr.Slider(
+            50, 500, value=150, step=25, label="Number of optimization steps"
+        )
+        control_variates_dropdown = gr.Dropdown(
+            choices=["None", "Lower", "Upper"],
+            value="None",
+            label="Control Variates",
+            info="Select which control variates to use for optimization",
+        )
+        candidates_per_pass_slider = gr.Slider(
+            0,
+            64,
+            value=16,
+            step=1,
+            label="Number of new candidates per pass.",
+            info=(
+                "The number of new candidates to generate per pass in the reservoir "
+                "sampling."
+            ),
+        )
+        num_projections_slider = gr.Slider(
+            16, 1024, value=64, step=16, label="Number of projections"
+        )
+        sampling_mode_dropdown = gr.Dropdown(
+            choices=["gaussian", "qmc"],
+            value="gaussian",
+            label="Sampling Mode",
+            info="Select which sampling mode to use for projections",
+        )
+
+    run_button = gr.Button("Run Color Matching", variant="primary")
+
+    with gr.Column(variant="compact"):
+        output_image = gr.Image(label="Color Matched Image", height=512)
+        output_cdl = gr.Textbox(label="CDL Parameters")
+
+    @spaces.GPU
+    def run_color_matching(
+        match_resolution: int,
+        num_steps: int,
+        control_variates: str,
+        num_projections: int,
+        candidates_per_pass: int,
+        sampling_mode: str,
+        learning_rate: float,
+        *images,
+    ):
+        """
+        Runs the color matching process between source and target images.
+        """
+        # Split images into source and target pair
+        source_img = images[0] / 255.0
+        target_img = images[1] / 255.0
+
+        # Convert images to tensors
+        source_tensors = to_torch(source_img).float().unsqueeze(0)
+        target_tensors = to_torch(target_img).float().unsqueeze(0)
+
+        # Configure loss function
+        criteria = VectorSWDLoss(
+            num_proj=num_projections,
+            use_ucv=control_variates == "Upper",
+            use_lcv=control_variates == "Lower",
+            num_new_candidates=candidates_per_pass,
+            sampling_mode=sampling_mode,
+        )
+
+        # Run color matching
+        source_matched, cdl, losses = train(
+            fabric=fabric,
+            criteria=criteria,
+            source_img=source_tensors,
+            target_img=target_tensors,
+            num_steps=num_steps,
+            lr=learning_rate,
+            match_resolution=match_resolution,
+        )
+
+        return [
+            (from_torch(source_matched.squeeze(0)) * 255).astype(np.uint8),
+            cdl.to_cdl_xml()[0],
+        ]
+
+    run_button.click(
+        run_color_matching,
+        inputs=[
+            match_resolution_slider,
+            num_steps_slider,
+            control_variates_dropdown,
+            num_projections_slider,
+            candidates_per_pass_slider,
+            sampling_mode_dropdown,
+            learning_rate_slider,
+            source_image,
+            target_image,
+        ],
+        outputs=[output_image, output_cdl],
+    )
+
+    clear_button = gr.ClearButton(
+        [source_image, target_image, output_image, output_cdl]
+    )
+
+    clear_button.click(lambda: None)

src/gradio_demo/sw_guidance.py

@@ -0,0 +1,336 @@
+import os
+from typing import Optional
+
+import gradio as gr
+import lightning as L
+import numpy as np
+import spaces
+from PIL import Image
+
+from src.sw_sdthree_guidance import create_pipeline
+from src.sw_sdthree_guidance import run as sw_guidance_run
+
+preset_lrs = [1e-6, 1.0]
+
+log_lrs = np.log10(preset_lrs)
+
+models = {
+    "stabilityai/stable-diffusion-3.5-large": {
+        "num_inference_steps": 30,
+        "guidance_scale": 7.0,
+        "sw_u_lr": np.log10(3.2e-3),
+        "sw_steps": 6,
+        "cfg_rescale_phi": 0.7,
+    },
+    "stabilityai/stable-diffusion-3.5-medium": {
+        "num_inference_steps": 30,
+        "guidance_scale": 3.5,
+        "sw_u_lr": np.log10(3.2e-3),
+        "sw_steps": 6,
+        "cfg_rescale_phi": 0.7,
+    },
+    "stabilityai/stable-diffusion-3.5-large-turbo": {
+        "num_inference_steps": 4,
+        "guidance_scale": 1.0,
+        "sw_u_lr": np.log10(4e-3),
+        "sw_steps": 6,
+        "cfg_rescale_phi": 0.65,
+    },
+}
+
+
+def log_slider_to_lr(log_lr):
+    return float(f"{10**log_lr:.1e}")
+
+
+def create_sw_guidance(
+    fabric: L.Fabric, model_name: str = "stabilityai/stable-diffusion-3.5-large"
+):
+    """
+    Creates the Gradio interface for SW guidance with SD3.5.
+
+    Args:
+        fabric: Lightning Fabric instance
+        model_name: The model to use for guidance
+    """
+    gr.Markdown(
+        f"""
+        # SW Guidance with {model_name.split('/')[-1]}
+        Generates images using SW Guidance with a reference image and text prompt.
+        """
+    )
+
+    pipe = create_pipeline(
+        model_name,
+        device="cuda",
+        compile=True,
+    )
+
+    model_config = models[model_name]
+
+    @spaces.GPU
+    def run_sw_guidance(
+        num_inference_steps: int,
+        num_guided_steps_perc: float,
+        guidance_scale: float,
+        sw_u_lr: float,
+        sw_steps: int,
+        num_projections: int,
+        control_variates: str,
+        distance: str,
+        candidates_per_pass: int,
+        subsampling_factor: int,
+        sampling_mode: str,
+        cfg_rescale_phi: float,
+        prompt: str,
+        reference_image: np.ndarray,
+        seed: Optional[int] = None,
+    ):
+        """
+        Runs the SW guidance process with the given parameters.
+        """
+        if reference_image is None:
+            raise gr.Error("Please provide a reference image")
+        if not prompt:
+            raise gr.Error("Please provide a prompt")
+
+        # Convert numpy array to PIL Image
+        ref_img = Image.fromarray(reference_image)
+
+        # Run SW guidance
+        image = sw_guidance_run(
+            prompt=prompt,
+            reference_image=ref_img,
+            model_path=model_name,
+            num_inference_steps=num_inference_steps,
+            num_guided_steps=int(num_guided_steps_perc * num_inference_steps),
+            guidance_scale=guidance_scale,
+            sw_u_lr=log_slider_to_lr(sw_u_lr),
+            sw_steps=sw_steps,
+            height=1024,
+            width=1024,
+            device="cuda",
+            num_projections=num_projections,
+            use_ucv=control_variates == "Upper",
+            use_lcv=control_variates == "Lower",
+            distance=distance,
+            num_new_candidates=candidates_per_pass,
+            subsampling_factor=subsampling_factor,
+            sampling_mode=sampling_mode,
+            pipe=pipe,
+            compile=True,
+            seed=seed,
+            cfg_rescale_phi=cfg_rescale_phi,
+        )
+
+        return np.array(image)
+
+    gr.Markdown("## Input")
+    with gr.Row(equal_height=True):
+        with gr.Column(variant="panel"):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your prompt here...",
+                lines=2,
+            )
+            reference_image = gr.Image(label="Reference Image", height=512)
+
+        with gr.Column(variant="panel"):
+            output_image = gr.Image(label="Generated Image", height=512)
+
+    example_pairs = [
+        ("A diver discovering an underwater city", "sunken_boat.jpg"),
+        ("A raccoon in a forest", "waterfall.jpg"),
+        ("A cat detective solving a mystery", "building.jpg"),
+        ("A raccoon reading a book by candlelight.", "food_stand.jpg"),
+        ("A lion meditating on a mountain", "mountain.jpg"),
+        ("A squirrel kayaking", "lake_green.jpg"),
+        ("A young dragon roasting marshmallows", "canyon.jpg"),
+        ("A family picnic beneath floating lanterns", "lake_sunset.jpg"),
+        ("Elephants holding umbrellas in a rainstorm", "fisher.jpg"),
+        ("A boy and his dog exploring a crystal cave", "lake.jpg"),
+        ("A snowman sharing cocoa with woodland animals", "snow.jpg"),
+        ("A kitten exploring an antique library", "ornament.jpg"),
+        ("Children riding flying bicycles over mountains", "sky_pier.jpg"),
+        ("An ancient tree whispering stories to deer", "greenhouse_2.jpg"),
+        ("Owls with monocles in treetops", "path.jpg"),
+    ]
+
+    default_config = {
+        "num_guided_steps_perc": 0.95,
+        "num_projections": 32,
+        "control_variates": "None",
+        "distance": "l1",
+        "candidates_per_pass": 8,
+        "subsampling_factor": 1,
+        "sampling_mode": "gaussian",
+    } | model_config
+
+    def run_example(prompt: str, reference_image: np.ndarray):
+        return run_sw_guidance(
+            **default_config,
+            prompt=prompt,
+            reference_image=reference_image,
+        )
+
+    example_inputs = [
+        [prompt, os.path.join("example", "guidance", img_file)]
+        for prompt, img_file in example_pairs
+    ]
+
+    run_button = gr.Button("Generate Image", variant="primary")
+
+    gr.Examples(
+        examples=example_inputs,
+        inputs=[prompt, reference_image],
+        outputs=[output_image],
+        fn=run_example,
+        label="Prompt + Reference Image Examples",
+        examples_per_page=5,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    gr.Markdown(
+        """
+        # Configuration
+        Adjust the parameters for the SW guidance process.
+        """
+    )
+
+    with gr.Accordion("Basic Config", open=True):
+        num_inference_steps_slider = gr.Slider(
+            1,
+            100,
+            value=model_config["num_inference_steps"],
+            step=1,
+            label="Number of inference steps",
+        )
+
+        guidance_scale_slider = gr.Slider(
+            0.0,
+            20.0,
+            value=model_config["guidance_scale"],
+            step=0.1,
+            label="Guidance scale",
+        )
+
+        cfg_rescale_phi_slider = gr.Slider(
+            0.0,
+            1.0,
+            value=model_config["cfg_rescale_phi"],
+            step=0.05,
+            label="CFG Rescale Phi",
+            info="Controls the rescaling of classifier-free guidance",
+        )
+
+        seed_input = gr.Number(
+            value=lambda: None,
+            label="Seed (leave empty for random)",
+            precision=0,
+            interactive=True,
+        )
+
+        with gr.Row(variant="panel", equal_height=True):
+            sw_u_lr_slider = gr.Slider(
+                minimum=log_lrs.min(),
+                maximum=log_lrs.max(),
+                value=model_config["sw_u_lr"],
+                step=0.05,
+                label="SW guidance learning rate (Log scale)",
+                interactive=True,
+                scale=4,
+            )
+            lr_display = gr.Textbox(
+                label="Learning Rate",
+                value=f"{log_slider_to_lr(model_config['sw_u_lr']):.1e}",
+                interactive=False,
+                scale=1,
+            )
+        sw_u_lr_slider.change(
+            lambda x: gr.update(value=f"{log_slider_to_lr(x):.1e}"),
+            inputs=sw_u_lr_slider,
+            outputs=lr_display,
+            show_progress=False,
+        )
+
+    with gr.Accordion("Advanced Config", open=False):
+        num_guided_steps_perc_slider = gr.Slider(
+            0.0,
+            1.0,
+            value=default_config["num_guided_steps_perc"],
+            step=0.05,
+            label="Percentage of steps to apply SW guidance",
+        )
+        sw_steps_slider = gr.Slider(
+            0,
+            32,
+            value=model_config["sw_steps"],
+            step=1,
+            label="Number of SW guidance steps, 0 means no SW guidance",
+        )
+        num_projections_slider = gr.Slider(
+            16,
+            1024,
+            value=default_config["num_projections"],
+            step=16,
+            label="Number of projections",
+        )
+        control_variates_dropdown = gr.Dropdown(
+            choices=["None", "Lower", "Upper"],
+            value=default_config["control_variates"],
+            label="Control Variates",
+            info="Select which control variates to use for optimization",
+        )
+        distance_dropdown = gr.Dropdown(
+            choices=["l1", "l2"],
+            value=default_config["distance"],
+            label="Distance metric",
+            info="Select which distance metric to use",
+        )
+        candidates_per_pass_slider = gr.Slider(
+            0,
+            64,
+            value=default_config["candidates_per_pass"],
+            step=1,
+            label="Number of new candidates per pass. 0 means no reservoir sampling",
+        )
+        subsampling_factor_slider = gr.Slider(
+            1,
+            16,
+            value=default_config["subsampling_factor"],
+            step=1,
+            label="Subsampling factor",
+        )
+        sampling_mode_dropdown = gr.Dropdown(
+            choices=["gaussian", "qmc"],
+            value="qmc",
+            label="Sampling Mode",
+            info="Select which sampling mode to use for projections",
+        )
+
+    run_button.click(
+        run_sw_guidance,
+        inputs=[
+            num_inference_steps_slider,
+            num_guided_steps_perc_slider,
+            guidance_scale_slider,
+            sw_u_lr_slider,
+            sw_steps_slider,
+            num_projections_slider,
+            control_variates_dropdown,
+            distance_dropdown,
+            candidates_per_pass_slider,
+            subsampling_factor_slider,
+            sampling_mode_dropdown,
+            cfg_rescale_phi_slider,
+            prompt,
+            reference_image,
+            seed_input,
+        ],
+        outputs=[output_image],
+    )
+
+    clear_button = gr.ClearButton([prompt, reference_image, output_image, seed_input])
+
+    clear_button.click(lambda: None)

src/loss/__init__.py

@@ -0,0 +1,4 @@
+from .abstract_loss import AbstractLoss
+from .vector_swd import VectorSWDLoss
+
+__all__ = ["AbstractLoss", "VectorSWDLoss"]

src/loss/abstract_loss.py

@@ -0,0 +1,20 @@
+import abc
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+
+
+class AbstractLoss(nn.Module, abc.ABC):
+    @abc.abstractmethod
+    def forward(
+        self,
+        pred: Float[torch.Tensor, "B C H W"],
+        gt: Float[torch.Tensor, "B C H W"],
+        step: int,
+        **kwargs,
+    ) -> Float[torch.Tensor, ""]:
+        pass
+
+    def reset(self):
+        pass

src/loss/vector_swd.py

@@ -0,0 +1,431 @@
+import typing
+
+import torch
+import torch.nn.functional as F
+from jaxtyping import Float
+
+from src.loss.abstract_loss import AbstractLoss
+from src.utils.math import sobol_sphere
+
+
+def process_vector(
+    x: Float[torch.Tensor, "B D N"],
+    dirs: Float[torch.Tensor, "K D"],
+) -> Float[torch.Tensor, "K B*N_valid"]:
+    """
+    Project a 1-D sequence with a bank of linear directions.
+
+    Args
+    ----
+    x      : (B, D, N) tensor – predictions or ground truth
+    dirs   : (K, D)   tensor – unit-length projection directions
+
+    Returns
+    -------
+    proj   : (K, B*N_valid) tensor of flattened projections
+    """
+    B, D, N = x.shape
+    K, _ = dirs.shape
+
+    # linear projection:   x  (B,D,N)  ->  (B,N,K)  ->  (K,B*N)
+    proj = F.linear(x.transpose(1, 2).to(torch.float32), dirs.to(torch.float32))
+    proj = proj.permute(2, 0, 1).reshape(K, -1).to(x.dtype)
+
+    return proj
+
+
+class VectorSWDLoss(AbstractLoss):
+    """
+    1-D Sliced-Wasserstein Distance on sequences.
+
+    This loss computes the sliced Wasserstein distance between predicted and ground
+    truth sequences by projecting them onto random directions and computing the
+    Wasserstein distance in 1D. It supports reservoir sampling for adaptive direction
+    selection and various variance reduction techniques.
+
+    Parameters
+    ----------
+    num_proj : int, default=64
+        Number of random projections to use per step (K).
+
+    distance : {"l1", "l2"}, default="l1"
+        Distance metric to use for computing the Wasserstein distance.
+
+    use_ucv : bool, default=False
+        Whether to use upper bounds control variates for variance reduction.
+        Mutually exclusive with use_lcv.
+
+    use_lcv : bool, default=False
+        Whether to use lower bounds control variates for variance reduction.
+        Mutually exclusive with use_ucv.
+
+    refresh_projections_every_n_steps : int, default=1
+        How often to refresh the projection directions. A value of 1 means
+        refresh every step, higher values reuse directions for multiple steps.
+
+    num_new_candidates : int, default=16
+        Number of new candidate directions to generate per step (M).
+        If 0, reservoir sampling is disabled. Must not exceed num_proj.
+
+    ess_alpha : float, default=0.5
+        Effective sample size threshold for resetting the reservoir.
+        When ESS drops below ess_alpha * reservoir_size, the reservoir is reset.
+
+    time_decay_tau : float or None, default=30.0
+        Time decay parameter for reservoir weights. If None, no time decay is applied.
+        Weights decay exponentially with age: exp(-age / time_decay_tau).
+
+    missing_value_method : {"random_replicate", "interpolate"},
+            default="random_replicate"
+        Method for handling sequences of different lengths:
+        - "random_replicate": Randomly replicate shorter sequences
+        - "interpolate": Use linear interpolation to match lengths
+
+    sampling_mode : {"gaussian", "qmc"}, default="qmc"
+        Method for generating random projection directions:
+        - "gaussian": Standard Gaussian sampling
+        - "qmc": Quasi-Monte Carlo sampling using Sobol sequences
+
+    Notes
+    -----
+    - Reservoir sampling is enabled when num_new_candidates > 0
+    - Reservoir size = num_proj - num_new_candidates
+    - When use_ucv or use_lcv is True, variance reduction is applied using
+      control variates based on the difference between sample and population means
+    - The loss automatically handles sequences of different lengths using the
+      specified missing_value_method
+    """
+
+    def __init__(
+        self,
+        num_proj: int = 64,
+        distance: typing.Literal["l1", "l2"] = "l1",
+        use_ucv: bool = False,
+        use_lcv: bool = False,
+        refresh_projections_every_n_steps: int = 1,
+        num_new_candidates: int = 16,
+        ess_alpha: float = 0.5,
+        time_decay_tau: float | None = 30.0,
+        missing_value_method: typing.Literal[
+            "random_replicate", "interpolate"
+        ] = "random_replicate",
+        sampling_mode: typing.Literal[
+            "gaussian",
+            "qmc",
+        ] = "qmc",
+    ):
+        super().__init__()
+
+        assert not (use_ucv and use_lcv), "use_ucv and use_lcv cannot both be True"
+
+        self.num_proj = num_proj
+        self.distance = distance
+        self.use_ucv = use_ucv
+        self.use_lcv = use_lcv
+
+        self.refresh_projections_every_n_steps = refresh_projections_every_n_steps
+        self.num_new_candidates = num_new_candidates  # M
+        self.ess_alpha = ess_alpha
+        self.time_decay_tau = time_decay_tau
+        self.missing_value_method = missing_value_method
+
+        if num_new_candidates > 0 and self.refresh_projections_every_n_steps != 1:
+            # Print a warning that this is not recommended
+            print(
+                "WARNING: num_new_candidates > 0 (enabling reservoir sampling) and "
+                "refresh_projections_every_n_steps != 1 is not recommended"
+            )
+        assert (
+            num_new_candidates <= num_proj
+        ), "`num_new_candidates` must not exceed `num_proj`"
+
+        # internal state for reservoir sampling
+        self.restir_enabled = self.num_new_candidates > 0
+        self.reservoir_size = self.num_proj - self.num_new_candidates
+        self.register_buffer("_reservoir_filters", torch.empty(0))
+        self.register_buffer("_reservoir_weights", torch.empty(0))
+        self.register_buffer("_reservoir_steps", torch.empty(0, dtype=torch.long))
+        self.register_buffer("_reservoir_keys", torch.empty(0))
+        self.register_buffer("_cumulative_weights", torch.tensor(0.0))
+        self.register_buffer("_has_reservoir", torch.tensor(False, dtype=torch.bool))
+
+        self._cached_dirs: typing.Optional[torch.Tensor] = None
+        self.sampling_mode = sampling_mode
+        self.sobol_engine = None
+
+    def _gaussian_proposals(self, k: int, d: int, device: torch.device) -> torch.Tensor:
+        """Generate Gaussian random projection directions."""
+        w = torch.randn(k, d, device=device)
+        return w / (w.norm(dim=1, keepdim=True) + 1e-8)  # unit length
+
+    def _qmc_proposals(self, k: int, d: int, device: torch.device) -> torch.Tensor:
+        """Generate quasi-Monte Carlo projection directions using Sobol sequences."""
+        vecs, self.sobol_engine = sobol_sphere(k, d, device, self.sobol_engine)
+        return vecs.view(k, d)
+
+    def _draw_dirs(self, k: int, d: int, device: torch.device) -> torch.Tensor:
+        """Draw projection directions using the specified sampling mode."""
+        if self.sampling_mode == "gaussian":
+            return self._gaussian_proposals(k, d, device)
+        if self.sampling_mode == "qmc":
+            return self._qmc_proposals(k, d, device)
+        raise ValueError("bad sampling_mode")
+
+    @staticmethod
+    def _duplicate_to_match(a: torch.Tensor, b: torch.Tensor, method: str):
+        """
+        Make two tensors have the same length by duplicating the shorter one.
+
+        Args
+        ----
+        a, b : (K, N₁) and (K, N₂) tensors
+        method : "random_replicate" or "interpolate"
+
+        Returns
+        -------
+        a, b : Tensors with matching second dimension
+        """
+        if a.shape[1] == b.shape[1]:
+            return a, b
+        if a.shape[1] < b.shape[1]:
+            a, b = b, a  # swap so that `a` is the larger
+
+        K, NA = a.shape
+        NB = b.shape[1]
+
+        # repeat / interpolate B until it matches A
+        if method == "random_replicate":
+            repeats = NA // NB
+            b = torch.cat([b] * repeats, dim=1)
+            if b.shape[1] < NA:
+                idx = torch.randint(0, NB, (NA - b.shape[1],), device=b.device)
+                b = torch.cat([b, b[:, idx]], dim=1)
+        else:  # interpolate
+            b = F.interpolate(
+                b.unsqueeze(0), size=(NA,), mode="linear", align_corners=False
+            ).squeeze(0)
+        return a, b
+
+    def reset(self):
+        """Reset the reservoir sampling state."""
+        if self.restir_enabled:
+            self._reservoir_filters = torch.empty(0)
+            self._reservoir_weights = torch.empty(0)
+            self._cumulative_weights.data.fill_(0)
+            self._has_reservoir.fill_(False)
+            self._reservoir_steps = torch.empty(0, dtype=torch.long)
+            self._reservoir_keys = torch.empty(0)
+
+    def _wrs_multi(
+        self, filters: torch.Tensor, weights: torch.Tensor, step: int
+    ) -> torch.Tensor:
+        """
+        Weighted reservoir sampling that keeps exactly self.reservoir_size samples and
+        returns their indices inside the concatenated candidate set.
+
+        Args
+        ----
+        filters : (K+M, D) tensor of candidate directions
+        weights : (K+M,) tensor of importance weights
+        step : Current training step
+
+        Returns
+        -------
+        keep_idx : Indices of kept samples
+        keep_w : Normalized weights of kept samples
+        """
+        R = self.reservoir_size
+        device = weights.device
+
+        u = torch.rand_like(weights)
+        keys = u.pow(1.0 / weights.clamp_min(1e-9))
+
+        if not self._has_reservoir.item():
+            self._reservoir_filters = filters[:R]
+            self._reservoir_weights = weights[:R]
+            self._reservoir_keys = keys[:R]
+            self._reservoir_steps = torch.full(
+                (R,), step, dtype=torch.long, device=device
+            )
+            self._has_reservoir.fill_(True)
+
+        new_filters = filters[R:]
+        new_keys = keys[R:]
+        new_weights = weights[R:]
+        new_steps = torch.full(
+            (new_filters.size(0),), step, dtype=torch.long, device=device
+        )
+
+        all_filters = torch.cat([self._reservoir_filters, new_filters], 0)
+        all_keys = torch.cat([self._reservoir_keys, new_keys], 0)
+        all_weights = torch.cat([self._reservoir_weights, new_weights], 0)
+        all_steps = torch.cat([self._reservoir_steps, new_steps], 0)
+
+        topk_keys, topk_idx = torch.topk(all_keys, R, largest=True)
+
+        self._reservoir_filters = all_filters[topk_idx]
+        self._reservoir_weights = all_weights[topk_idx]
+        self._reservoir_keys = topk_keys
+        self._reservoir_steps = all_steps[topk_idx]
+
+        # indices w.r.t. current cand_dirs (old R first, then new M)
+        keep_idx = torch.cat(
+            [
+                torch.arange(R, device=device),
+                torch.arange(R, R + new_filters.size(0), device=device),
+            ]
+        )[topk_idx]
+        keep_w = self._reservoir_weights / self._reservoir_weights.sum().clamp_min(
+            1e-12
+        )
+        return keep_idx, keep_w
+
+    def _apply_time_decay(self, step: int):
+        """
+        Apply exponential time decay to stored reservoir weights.
+
+        Args
+        ----
+        step : Current training step
+        """
+        if self.time_decay_tau is None or not self._has_reservoir.item():
+            return
+        age = (step - self._reservoir_steps).to(torch.float32)
+        decay = torch.exp(-age / self.time_decay_tau).to(self._reservoir_weights.dtype)
+        self._reservoir_weights.mul_(decay)
+        self._reservoir_keys.mul_(decay)  # preserve ordering consistency
+
+    def forward(
+        self,
+        pred: Float[torch.Tensor, "B D N"],
+        gt: Float[torch.Tensor, "B D N"],
+        step: int,
+    ):
+        """
+        Compute the sliced Wasserstein distance between predicted and ground truth
+        sequences.
+
+        Args
+        ----
+        pred : (B, D, N) tensor of predicted sequences
+        gt : (B, D, N) tensor of ground truth sequences
+        step : Current training step for reservoir sampling
+
+        Returns
+        -------
+        loss : Scalar tensor containing the computed loss
+        """
+        B, D, N = pred.shape
+        K = self.num_proj
+        M = self.num_new_candidates
+        R = self.reservoir_size
+        device = pred.device
+        gt = gt.detach()
+
+        self._apply_time_decay(step)
+
+        # Get candidate directions
+        if step % self.refresh_projections_every_n_steps == 0:
+            new_dirs = self._draw_dirs(
+                M if self.restir_enabled and self._has_reservoir.item() else K,
+                D,
+                device,
+            )
+            self._cached_dirs = new_dirs
+        else:
+            new_dirs = self._cached_dirs
+
+        if self.restir_enabled and self._has_reservoir.item():
+            cand_dirs = torch.cat(
+                [self._reservoir_filters, new_dirs], dim=0
+            )  # [K+M, C,P,P]
+        else:
+            cand_dirs = new_dirs
+
+        # Project sequences
+        cand_pred = process_vector(pred, cand_dirs)
+        cand_gt = process_vector(gt, cand_dirs)
+
+        cand_pred, cand_gt = self._duplicate_to_match(
+            cand_pred, cand_gt, self.missing_value_method
+        )
+
+        cand_pred = cand_pred.sort(dim=1).values
+        cand_gt = cand_gt.sort(dim=1).values
+
+        # Select K directions (reservoir) & importance weights
+        if self.restir_enabled:
+            with torch.no_grad():
+                base = cand_pred - cand_gt
+                base = base.abs() if self.distance == "l1" else base.square()
+                ris_weights = base.mean(1)  # (K+M)
+                keep_idx, keep_w = self._wrs_multi(cand_dirs, ris_weights, step)
+
+            w = keep_w
+            w_hat = keep_w
+
+            dirs = cand_dirs[keep_idx]
+            proj_pred = cand_pred[keep_idx]
+            proj_gt = cand_gt[keep_idx]
+        else:
+            dirs = cand_dirs
+            proj_pred = cand_pred
+            proj_gt = cand_gt
+            w = torch.full((dirs.shape[0],), 1.0 / K, device=device)
+
+        # Compute SWD
+        diff = proj_pred - proj_gt
+        diff = diff.abs() if self.distance == "l1" else diff.square()
+        per_slice = diff.mean(1)  # (L,)
+
+        if self.use_ucv or self.use_lcv:
+            X_vecs = pred.permute(0, 2, 1).reshape(-1, D)  # (B·N, D)
+            Y_vecs = gt.permute(0, 2, 1).reshape(-1, D)  # (B·N, D)
+
+            m1 = X_vecs.mean(0)  # (D,)
+            m2 = Y_vecs.mean(0)
+            diff_m = m1 - m2  # (D,)
+
+            theta = dirs  # (L, D) already unit-norm
+
+            if self.use_ucv:
+                diff_X = X_vecs - m1
+                diff_Y = Y_vecs - m2
+
+                d = D
+                trSigX = diff_X.pow(2).mean()
+                trSigY = diff_Y.pow(2).mean()
+                G_bar = (diff_m @ diff_m) / d + (trSigX + trSigY)
+
+                delta2 = (theta @ diff_m) ** 2  # (L,)
+
+                proj_X = diff_X @ theta.t()  # (B·N, L)
+                proj_Y = diff_Y @ theta.t()
+                varX = proj_X.pow(2).mean(0)  # (L,)
+                varY = proj_Y.pow(2).mean(0)
+                G_hat = delta2 + varX + varY
+            else:  # LCV
+                d = D
+                G_bar = (diff_m @ diff_m) / d
+                G_hat = (theta @ diff_m) ** 2
+
+            diff_hat_G_mean_G = G_hat - G_bar
+
+            hat_A = (w * per_slice).sum()
+            var_G = (w * diff_hat_G_mean_G.pow(2)).sum()
+            cov_AG = (w * (per_slice - hat_A) * diff_hat_G_mean_G).sum()
+            hat_alpha = cov_AG / (var_G + 1e-12)
+            loss = hat_A - hat_alpha * (w * diff_hat_G_mean_G).sum()
+        else:
+            loss = (w * per_slice).sum()
+
+        # Reservoir update
+        if self.restir_enabled and self.ess_alpha > 0:
+            with torch.no_grad():
+                ess = (w_hat.sum().square()) / (w_hat.square().sum() + 1e-12)
+                ess = torch.nan_to_num(ess, nan=0.0, posinf=R, neginf=0.0).item()
+                if ess < self.ess_alpha * R:
+                    print(f"ESS: {ess} is less than {self.ess_alpha * R}, resetting")
+                    self.reset()
+
+        return loss

src/sw_sdthree_guidance.py

@@ -0,0 +1,654 @@
+# Implementation of the SW Guidance method with our enhanced SWD implementation
+# See: https://github.com/alobashev/sw-guidance/ for the original implementation
+#
+# Alexander Lobashev, Maria Larchenko, Dmitry Guskov
+# Color Conditional Generation with Sliced Wasserstein Guidance
+# https://arxiv.org/abs/2503.19034
+
+
+import gc
+import os
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
+
+import numpy as np
+import PIL
+import torch
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    StableDiffusion3Pipeline,
+)
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (
+    XLA_AVAILABLE,
+    StableDiffusion3PipelineOutput,
+    calculate_shift,
+    retrieve_timesteps,
+)
+
+from src.loss.vector_swd import VectorSWDLoss
+from src.utils.color_space import rgb_to_lab
+from src.utils.image import from_torch, write_img
+
+if XLA_AVAILABLE:
+    from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import xm
+
+
+def _no_grad_noise(model, *args, **kw):
+    """Forward pass with grad disabled; result is returned detached."""
+    with torch.no_grad():
+        return model(*args, **kw, return_dict=False)[0].detach()
+
+
+# ---------------- explicit pipeline forward call
+class SWStableDiffusion3Pipeline(StableDiffusion3Pipeline):
+    swd: VectorSWDLoss = None
+
+    def setup_swd(
+        self,
+        num_projections: int = 64,
+        use_ucv: bool = False,
+        use_lcv: bool = False,
+        distance: Literal["l1", "l2"] = "l1",
+        num_new_candidates: int = 32,
+        subsampling_factor: int = 1,
+        sampling_mode: Literal["gaussian", "qmc"] = "qmc",
+    ):
+        self.swd = VectorSWDLoss(
+            num_proj=num_projections,
+            distance=distance,
+            use_ucv=use_ucv,
+            use_lcv=use_lcv,
+            num_new_candidates=num_new_candidates,
+            missing_value_method="interpolate",
+            ess_alpha=-1,
+            sampling_mode=sampling_mode,
+        ).to(self.device)
+        self.subsampling_factor = subsampling_factor
+
+    def do_sw_guidance(
+        self,
+        sw_steps,
+        sw_u_lr,
+        latents,
+        t,
+        prompt_embeds,
+        pooled_prompt_embeds,
+        pixels_ref,
+        cur_iter_step,
+        write_video_animation_path: Optional[str] = None,
+    ):
+        if sw_steps == 0:
+            return latents
+
+        if latents.shape[0] != prompt_embeds.shape[0]:
+            prompt_embeds = prompt_embeds[1].unsqueeze(0)
+            pooled_prompt_embeds = pooled_prompt_embeds[1].unsqueeze(0)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timestep = t.expand(latents.shape[0])
+
+        pixels_ref = (
+            rgb_to_lab(pixels_ref.unsqueeze(0).clamp(0, 1).permute(0, 3, 1, 2))
+            .permute(0, 2, 3, 1)
+            .contiguous()
+        )
+
+        csc_scaler = torch.tensor(
+            [100, 2 * 128, 2 * 128], dtype=torch.bfloat16, device=latents.device
+        ).view(1, 3, 1)
+        csc_bias = torch.tensor(
+            [0, 0.5, 0.5], dtype=torch.bfloat16, device=latents.device
+        ).view(1, 3, 1)
+
+        u = torch.nn.Parameter(
+            torch.zeros_like(latents, dtype=torch.bfloat16, device=latents.device)
+        )
+        optimizer = torch.optim.Adam([u], lr=sw_u_lr)
+
+        for tt in range(sw_steps):
+            optimizer.zero_grad()
+
+            x_hat_t = latents.detach() + u
+            noise_pred = _no_grad_noise(
+                self.transformer,
+                hidden_states=x_hat_t,
+                timestep=timestep,
+                encoder_hidden_states=prompt_embeds,
+                pooled_projections=pooled_prompt_embeds,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+            )
+
+            # ------------ Compute x_0
+            sigma_t = self.scheduler.sigmas[
+                self.scheduler.index_for_timestep(t)
+            ]  # scalar
+            while sigma_t.ndim < x_hat_t.ndim:
+                sigma_t = sigma_t.unsqueeze(-1)
+            sigma_t = sigma_t.to(x_hat_t.dtype).to(latents.device)
+
+            x_0 = x_hat_t - sigma_t * noise_pred
+
+            # ------------ Compute loss
+            img_unscaled = self.vae.decode(
+                (x_0 / self.vae.config.scaling_factor) + self.vae.config.shift_factor,
+                return_dict=False,
+            )[0]
+            image = (img_unscaled * 0.5 + 0.5).clamp(0, 1)
+            image_matched = (
+                rgb_to_lab(image.permute(0, 3, 1, 2)).permute(0, 2, 3, 1).contiguous()
+            )
+            # reshape to (B, D, N)  where D=3, N = H*W
+            pred_seq = image_matched.view(1, 3, -1) / csc_scaler + csc_bias
+            ref_seq = pixels_ref.view(1, 3, -1) / csc_scaler + csc_bias
+
+            # Apply subsampling if enabled
+            if self.subsampling_factor > 1:
+                pred_seq = pred_seq[..., :: self.subsampling_factor]
+                ref_seq = ref_seq[..., :: self.subsampling_factor]
+
+            loss = self.swd(pred=pred_seq, gt=ref_seq, step=tt)
+
+            loss.backward()
+            optimizer.step()
+
+            if write_video_animation_path is not None:
+                frame_idx = cur_iter_step * sw_steps + tt
+                write_img(
+                    os.path.join(write_video_animation_path, f"{frame_idx:05d}.jpg"),
+                    from_torch(img_unscaled.squeeze(0)),
+                )
+
+        latents = latents.detach() + u.detach()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        return latents
+
+    def __call__(
+        self,
+        sw_reference: PIL.Image = None,
+        sw_steps: int = 8,
+        sw_u_lr: float = 0.05 * 10**3,
+        num_guided_steps: int = None,
+        # -----------------------------------
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_3: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 7.0,
+        cfg_rescale_phi: float = 0.7,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 256,
+        skip_guidance_layers: List[int] = None,
+        skip_layer_guidance_scale: float = 2.8,
+        skip_layer_guidance_stop: float = 0.2,
+        skip_layer_guidance_start: float = 0.01,
+        mu: Optional[float] = None,
+        write_video_animation_path: Optional[str] = None,
+    ):
+        assert self.swd is not None, "SWD not initialized"
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._skip_layer_guidance_scale = skip_layer_guidance_scale
+        self._clip_skip = clip_skip
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None)
+            if self.joint_attention_kwargs is not None
+            else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            clip_skip=self.clip_skip,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            if skip_guidance_layers is not None:
+                original_prompt_embeds = prompt_embeds
+                original_pooled_prompt_embeds = pooled_prompt_embeds
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        scheduler_kwargs = {}
+        if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None:
+            _, _, height, width = latents.shape
+            image_seq_len = (height // self.transformer.config.patch_size) * (
+                width // self.transformer.config.patch_size
+            )
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.16),
+            )
+            scheduler_kwargs["mu"] = mu
+        elif mu is not None:
+            scheduler_kwargs["mu"] = mu
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            **scheduler_kwargs,
+        )
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare image embeddings
+        if (
+            ip_adapter_image is not None and self.is_ip_adapter_active
+        ) or ip_adapter_image_embeds is not None:
+            ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+            if self.joint_attention_kwargs is None:
+                self._joint_attention_kwargs = {
+                    "ip_adapter_image_embeds": ip_adapter_image_embeds
+                }
+            else:
+                self._joint_attention_kwargs.update(
+                    ip_adapter_image_embeds=ip_adapter_image_embeds
+                )
+
+        if sw_reference is not None:
+            # Resize so the reference is maximal width or height of the output image
+
+            target_max_size = max(height, width)
+            reference_max_size = max(sw_reference.width, sw_reference.height)
+            scale_factor = target_max_size / reference_max_size
+
+            sw_reference = sw_reference.resize(
+                (
+                    int(sw_reference.width * scale_factor),
+                    int(sw_reference.height * scale_factor),
+                )
+            )
+            pixels_ref = (
+                torch.Tensor(np.array(sw_reference).astype(np.float32) / 255)
+                .permute(2, 0, 1)
+                .to(device)
+                .to(torch.bfloat16)
+            )
+
+        # 7. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # broadcast to batch dimension in a way that's compatible
+                # with ONNX/Core ML
+                timestep = t.expand(latents.shape[0])
+
+                # SW Guidance
+                if sw_reference is not None:
+                    if num_guided_steps is None or i < num_guided_steps:
+                        latents = self.do_sw_guidance(
+                            sw_steps,
+                            sw_u_lr,
+                            latents,
+                            t,
+                            prompt_embeds,
+                            pooled_prompt_embeds,
+                            pixels_ref,
+                            cur_iter_step=i,
+                            write_video_animation_path=write_video_animation_path,
+                        )
+                        if i == num_guided_steps // 2:
+                            self.swd.reset()
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2)
+                    if self.do_classifier_free_guidance
+                    else latents
+                )
+
+                with torch.no_grad():
+                    timestep = t.expand(latent_model_input.shape[0])
+
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        pooled_projections=pooled_prompt_embeds,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                    # perform guidance
+                    if self.do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (
+                            noise_pred_text - noise_pred_uncond
+                        )
+
+                        should_skip_layers = (
+                            True
+                            if i > num_inference_steps * skip_layer_guidance_start
+                            and i < num_inference_steps * skip_layer_guidance_stop
+                            else False
+                        )
+                        if skip_guidance_layers is not None and should_skip_layers:
+                            timestep = t.expand(latents.shape[0])
+                            latent_model_input = latents
+                            noise_pred_skip_layers = self.transformer(
+                                hidden_states=latent_model_input,
+                                timestep=timestep,
+                                encoder_hidden_states=original_prompt_embeds,
+                                pooled_projections=original_pooled_prompt_embeds,
+                                joint_attention_kwargs=self.joint_attention_kwargs,
+                                return_dict=False,
+                                skip_layers=skip_guidance_layers,
+                            )[0]
+                            noise_pred = (
+                                noise_pred
+                                + (noise_pred_text - noise_pred_skip_layers)
+                                * self._skip_layer_guidance_scale
+                            )
+
+                        # Based on Sec. 3.4 of Lin, Liu, Li, Yang -
+                        # Common Diffusion Noise Schedules and Sample Steps are Flawed
+                        # https://arxiv.org/abs/2305.08891
+                        # While Flow matching is free of most issues, a high CFG scale
+                        # can still cause over-exposure issues as discussed in the work.
+                        if cfg_rescale_phi is not None and cfg_rescale_phi > 0:
+                            # σ_pos and σ_cfg are per-sample (B×1×1×1) stdevs
+                            sigma_pos = noise_pred_text.std(dim=(1, 2, 3), keepdim=True)
+                            sigma_cfg = noise_pred.std(dim=(1, 2, 3), keepdim=True)
+
+                            # Linear blend between the raw ratio and 1,
+                            # cf. Eq. (15–16) in the paper
+                            factor = torch.lerp(
+                                sigma_pos / (sigma_cfg + 1e-8),  # avoid div-by-zero
+                                torch.ones_like(sigma_cfg),
+                                1.0 - cfg_rescale_phi,
+                            )
+                            noise_pred = noise_pred * factor
+                        else:
+                            noise_pred = noise_pred
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents_dtype = latents.dtype
+                    latents = self.scheduler.step(
+                        noise_pred, t, latents, return_dict=False
+                    )[0]
+
+                    if latents.dtype != latents_dtype:
+                        if torch.backends.mps.is_available():
+                            # some platforms (eg. apple mps) misbehave due to a
+                            # pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                            latents = latents.to(latents_dtype)
+
+                    if callback_on_step_end is not None:
+                        callback_kwargs = {}
+                        for k in callback_on_step_end_tensor_inputs:
+                            callback_kwargs[k] = locals()[k]
+                        callback_outputs = callback_on_step_end(
+                            self, i, t, callback_kwargs
+                        )
+
+                        latents = callback_outputs.pop("latents", latents)
+                        prompt_embeds = callback_outputs.pop(
+                            "prompt_embeds", prompt_embeds
+                        )
+                        negative_prompt_embeds = callback_outputs.pop(
+                            "negative_prompt_embeds", negative_prompt_embeds
+                        )
+                        negative_pooled_prompt_embeds = callback_outputs.pop(
+                            "negative_pooled_prompt_embeds",
+                            negative_pooled_prompt_embeds,
+                        )
+
+                    if write_video_animation_path is not None and i >= num_guided_steps:
+                        with torch.no_grad():
+                            image = self.vae.decode(
+                                (latents / self.vae.config.scaling_factor)
+                                + self.vae.config.shift_factor,
+                                return_dict=False,
+                            )[0]
+                            cur_frame_idx = i * sw_steps
+                            write_img(
+                                os.path.join(
+                                    write_video_animation_path,
+                                    f"{cur_frame_idx:05d}.jpg",
+                                ),
+                                from_torch(image.squeeze(0)),
+                            )
+
+                    # call the callback, if provided
+                    if i == len(timesteps) - 1 or (
+                        (i + 1) > num_warmup_steps
+                        and (i + 1) % self.scheduler.order == 0
+                    ):
+                        progress_bar.update()
+
+                    if XLA_AVAILABLE:
+                        xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (
+                latents / self.vae.config.scaling_factor
+            ) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(
+                image.detach(), output_type=output_type
+            )
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusion3PipelineOutput(images=image)
+
+
+def run(
+    prompt: str,
+    reference_image: PIL.Image.Image,
+    model_path: str,
+    num_inference_steps: int = 30,
+    num_guided_steps: int = 28,
+    guidance_scale: float = 5.0,
+    cfg_rescale_phi: float = 0.7,
+    sw_u_lr: float = 3e-3,
+    sw_steps: int = 8,
+    height: int = 768,
+    width: int = 768,
+    device: str = "cuda",
+    seed: Optional[int] = None,
+    # Add new SW-related parameters
+    num_projections: int = 64,
+    use_ucv: bool = False,
+    use_lcv: bool = False,
+    distance: Literal["l1", "l2"] = "l1",
+    num_new_candidates: int = 32,
+    subsampling_factor: int = 1,
+    sampling_mode: Literal["gaussian", "qmc"] = "gaussian",
+    pipe: Optional[SWStableDiffusion3Pipeline] = None,
+    compile: bool = False,
+    video_animation_path: Optional[str] = None,
+) -> PIL.Image.Image:
+    """
+    Generate an image using SW Guidance with a given prompt and reference image.
+
+    Args:
+        prompt (str): Text prompt to guide the generation
+        reference_image (PIL.Image.Image): Reference image to guide the generation
+        model_path (str): Path to the model weights
+        num_inference_steps (int): Number of denoising steps
+        num_guided_steps (int): Number of steps to apply SW guidance
+        guidance_scale (float): Scale for classifier-free guidance
+        cfg_rescale_phi (float): Rescale factor for classifier-free guidance
+        sw_u_lr (float): Learning rate for SW guidance
+        sw_steps (int): Number of steps to apply SW guidance
+        height (int): Output image height
+        width (int): Output image width
+        device (str): Device to run the model on
+        num_projections (int): Number of random projections for VectorSWDLoss
+        use_ucv (bool): Use UCV variant of VectorSWDLoss
+        use_lcv (bool): Use LCV variant of VectorSWDLoss
+        distance (str): Distance metric for VectorSWDLoss ("l1" or "l2")
+        refresh_projections_every_n_steps (int): How often to refresh projections
+        num_new_candidates (int): Number of new candidates for the reservoir
+        subsampling_factor (int): Factor to subsample points for SW computation.
+            Higher values reduce memory usage but may affect quality.
+        sampling_mode (str): Sampling mode for VectorSWDLoss.
+        pipe (SWStableDiffusion3Pipeline): Pipeline to use for generation.
+            If None, a new pipeline is created.
+        compile (bool): Whether to compile the pipeline.
+
+    Returns:
+        PIL.Image.Image: Generated image
+    """
+    # Normalize device to torch.device for robustness
+    device = torch.device(device) if not isinstance(device, torch.device) else device
+    if pipe is None:
+        pipe = create_pipeline(model_path, device, compile=compile)
+
+    pipe.setup_swd(
+        num_projections=num_projections,
+        use_ucv=use_ucv,
+        use_lcv=use_lcv,
+        distance=distance,
+        num_new_candidates=num_new_candidates,
+        subsampling_factor=subsampling_factor,
+        sampling_mode=sampling_mode,
+    )
+
+    if seed is not None:
+        print(f"Using seed: {seed}")
+        generator = torch.Generator(device=device).manual_seed(seed)
+    else:
+        generator = None
+
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        num_guided_steps=num_guided_steps,
+        guidance_scale=guidance_scale,
+        cfg_rescale_phi=cfg_rescale_phi,
+        sw_u_lr=sw_u_lr,
+        sw_steps=sw_steps,
+        height=height,
+        width=width,
+        sw_reference=reference_image,
+        generator=generator,
+        write_video_animation_path=video_animation_path,
+    ).images[0]
+
+    return image
+
+
+def create_pipeline(model_path, device: str = "cuda", compile: bool = False):
+    pipe = SWStableDiffusion3Pipeline.from_pretrained(
+        model_path,
+        torch_dtype=torch.bfloat16,
+    )
+    pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
+    if compile:
+        pipe.transformer = torch.compile(pipe.transformer)
+        pipe.vae.decoder = torch.compile(pipe.vae.decoder)
+    return pipe

src/utils/asc_cdl.py

@@ -0,0 +1,279 @@
+import warnings
+from typing import Optional
+
+import torch
+from jaxtyping import Float
+from lxml import etree
+
+
+def load_asc_cdl(cdl_path: str, device: torch.device = torch.device("cpu")) -> dict:
+    """
+    Loads ASC CDL parameters from an XML file.
+
+    Parameters:
+        cdl_path (str): Path to the ASC CDL XML file
+
+    Returns:
+        Dict:
+            slope, offset, power, and saturation values as torch tensors
+    """
+    try:
+        tree = etree.parse(cdl_path)
+        root = tree.getroot()
+    except Exception as e:
+        raise ValueError(f"Error loading ASC CDL from {cdl_path}: {e}")
+
+    # Extract SOP values
+    sop_node = root.find(".//SOPNode")
+    slope = torch.tensor(
+        [float(x) for x in sop_node.find("Slope").text.split()], device=device
+    )
+    offset = torch.tensor(
+        [float(x) for x in sop_node.find("Offset").text.split()], device=device
+    )
+    power = torch.tensor(
+        [float(x) for x in sop_node.find("Power").text.split()], device=device
+    )
+
+    # Extract Saturation value
+    sat_node = root.find(".//SatNode")
+    saturation = torch.tensor(float(sat_node.find("Saturation").text), device=device)
+
+    return {"slope": slope, "offset": offset, "power": power, "saturation": saturation}
+
+
+def save_asc_cdl(cdl_dict: dict, cdl_path: Optional[str]):
+    """
+    Saves ASC CDL parameters to an XML file.
+
+    Parameters:
+        cdl_dict (dict): Dictionary containing slope, offset, power, and
+            saturation values
+    """
+    root = etree.Element("ASC_CDL")
+    sop_node = etree.SubElement(root, "SOPNode")
+    etree.SubElement(sop_node, "Slope").text = " ".join(
+        str(x) for x in cdl_dict["slope"].detach().cpu().numpy()
+    )
+    etree.SubElement(sop_node, "Offset").text = " ".join(
+        str(x) for x in cdl_dict["offset"].detach().cpu().numpy()
+    )
+    etree.SubElement(sop_node, "Power").text = " ".join(
+        str(x) for x in cdl_dict["power"].detach().cpu().numpy()
+    )
+    sat_node = etree.SubElement(root, "SatNode")
+    etree.SubElement(sat_node, "Saturation").text = str(
+        cdl_dict["saturation"].detach().cpu().numpy()
+    )
+
+    tree = etree.ElementTree(root)
+    if cdl_path is not None:
+        try:
+            tree.write(
+                cdl_path, pretty_print=True, xml_declaration=True, encoding="utf-8"
+            )
+        except Exception as e:
+            raise ValueError(f"Error saving ASC CDL to {cdl_path}: {e}")
+    else:
+        return etree.tostring(
+            root, pretty_print=True, xml_declaration=True, encoding="utf-8"
+        ).decode("utf-8")
+
+
+def apply_sop(
+    img: Float[torch.Tensor, "*B C H W"],
+    slope: Float[torch.Tensor, "*B C"],
+    offset: Float[torch.Tensor, "*B C"],
+    power: Float[torch.Tensor, "*B C"],
+    clamp: bool = True,
+) -> Float[torch.Tensor, "*B C H W"]:
+    """
+    Applies Slope, Offset, and Power adjustments.
+
+    Parameters:
+        img (torch.Tensor): Input image tensor (*B, C, H, W)
+        slope (torch.Tensor): Slope per channel (*B, C)
+        offset (torch.Tensor): Offset per channel (*B, C)
+        power (torch.Tensor): Power per channel (*B, C)
+
+    Returns:
+        torch.Tensor: Image after SOP adjustments.
+    """
+    so = img * slope.unsqueeze(-1).unsqueeze(-1) + offset.unsqueeze(-1).unsqueeze(-1)
+    if clamp:
+        so = torch.clamp(so, min=0.0, max=1.0)
+    return torch.where(
+        so > 1e-7, torch.pow(so.clamp(min=1e-7), power.unsqueeze(-1).unsqueeze(-1)), so
+    )
+
+
+def apply_saturation(
+    img: Float[torch.Tensor, "*B C H W"],
+    saturation: Float[torch.Tensor, "*B"],
+) -> Float[torch.Tensor, "*B C H W"]:
+    """
+    Applies saturation adjustment.
+
+    Parameters:
+        img (torch.Tensor): Image tensor (*B, C, H, W)
+        saturation (torch.Tensor): Saturation factor (*B)
+
+    Returns:
+        torch.Tensor: Image after saturation adjustment.
+    """
+    # Calculate luminance using Rec. 709 coefficients
+    lum = (
+        0.2126 * img[..., 0, :, :]
+        + 0.7152 * img[..., 1, :, :]
+        + 0.0722 * img[..., 2, :, :]
+    )
+    lum = lum.unsqueeze(-3)  # Add channel dimension
+    return lum + (img - lum) * saturation.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+
+
+def asc_cdl_forward(
+    img: Float[torch.Tensor, "*B C H W"],
+    slope: Float[torch.Tensor, "*B C"],
+    offset: Float[torch.Tensor, "*B C"],
+    power: Float[torch.Tensor, "*B C"],
+    saturation: Float[torch.Tensor, "*B"],
+    clamp: bool = True,
+) -> Float[torch.Tensor, "*B C H W"]:
+    """
+    Applies ASC CDL transformation in Fwd or FwdNoClamp mode.
+
+    Parameters:
+        img (torch.Tensor): Input image tensor (*B, C, H, W)
+        slope (torch.Tensor): Slope per channel (*B, C)
+        offset (torch.Tensor): Offset per channel (*B, C)
+        power (torch.Tensor): Power per channel (*B, C)
+        saturation (torch.Tensor): Saturation factor (*B)
+        clamp (bool): If True, clamps output to [0, 1] (Fwd mode).
+            If False, no clamping (FwdNoClamp mode).
+
+    Returns:
+        torch.Tensor: Transformed image tensor.
+    """
+    # Add warning if saturation, slope, power are below 0
+    if (saturation < 0).any():
+        warnings.warn("Saturation is below 0, this will result in a color shift.")
+    if (slope < 0).any():
+        warnings.warn("Slope is below 0, this will result in a color shift.")
+    if (power < 0).any():
+        warnings.warn("Power is below 0, this will result in a color shift.")
+
+    img_batch_dim = img.shape[:-3]
+    # Check if slope, offset, power, saturation have the same batch dimension
+    # If they do not have any batch dimensions, add a single batch dimensions
+    if slope.ndim == 1:
+        slope = slope.view(*[1] * len(img_batch_dim), *slope.shape)
+    if offset.ndim == 1:
+        offset = offset.view(*[1] * len(img_batch_dim), *offset.shape)
+    if power.ndim == 1:
+        power = power.view(*[1] * len(img_batch_dim), *power.shape)
+    if saturation.ndim == 0:
+        saturation = saturation.view(*[1] * len(img_batch_dim), *saturation.shape)
+
+    # Now check that the lengths are matching
+    assert slope.ndim == len(img_batch_dim) + 1
+    assert offset.ndim == len(img_batch_dim) + 1
+    assert power.ndim == len(img_batch_dim) + 1
+    assert saturation.ndim == len(img_batch_dim)
+
+    # Apply Slope, Offset, and Power adjustments
+    img = apply_sop(img, slope, offset, power, clamp=clamp)
+    # print("img after sop", img.min(), img.max())
+    # Apply Saturation adjustment
+    img = apply_saturation(img, saturation)
+    # print("img after saturation", img.min(), img.max())
+    # Clamp if in Fwd mode
+    if clamp:
+        img = torch.clamp(img, 0.0, 1.0)
+    return img
+
+
+def inverse_saturation(
+    img: Float[torch.Tensor, "*B C H W"],
+    saturation: Float[torch.Tensor, "*B"],
+) -> Float[torch.Tensor, "*B C H W"]:
+    """
+    Reverts saturation adjustment.
+
+    Parameters:
+        img (torch.Tensor): Image tensor (*B, C, H, W)
+        saturation (torch.Tensor): Saturation factor (*B)
+
+    Returns:
+        torch.Tensor: Image after reversing saturation adjustment.
+    """
+    # Calculate luminance using Rec. 709 coefficients
+    lum = (
+        0.2126 * img[..., 0, :, :]
+        + 0.7152 * img[..., 1, :, :]
+        + 0.0722 * img[..., 2, :, :]
+    )
+    lum = lum.unsqueeze(-3)  # Add channel dimension
+    return lum + (img - lum) / saturation.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+
+
+def asc_cdl_reverse(
+    img: Float[torch.Tensor, "*B C H W"],
+    slope: Float[torch.Tensor, "*B C"],
+    offset: Float[torch.Tensor, "*B C"],
+    power: Float[torch.Tensor, "*B C"],
+    saturation: Float[torch.Tensor, "*B"],
+    clamp: bool = True,
+) -> Float[torch.Tensor, "*B C H W"]:
+    """
+    Applies reverse ASC CDL transformation.
+
+    Parameters:
+        img (torch.Tensor): Transformed image tensor (*B, C, H, W)
+        slope (torch.Tensor): Slope per channel (*B, C)
+        offset (torch.Tensor): Offset per channel (*B, C)
+        power (torch.Tensor): Power per channel (*B, C)
+        saturation (torch.Tensor): Saturation factor (*B)
+        clamp (bool): If True, clamps output to [0, 1].
+
+    Returns:
+        torch.Tensor: Recovered input image tensor.
+    """
+    # Add warning if saturation, slope, power are below 0
+    if (saturation < 0).any():
+        warnings.warn("Saturation is below 0, this will result in a color shift.")
+    if (slope < 0).any():
+        warnings.warn("Slope is below 0, this will result in a color shift.")
+    if (power < 0).any():
+        warnings.warn("Power is below 0, this will result in a color shift.")
+
+    img_batch_dim = img.shape[:-3]
+    # Check if slope, offset, power, saturation have the same batch dimension
+    # If they do not have any batch dimensions, add a single batch dimensions
+    if slope.ndim == 1:
+        slope = slope.view(*[1] * len(img_batch_dim), *slope.shape)
+    if offset.ndim == 1:
+        offset = offset.view(*[1] * len(img_batch_dim), *offset.shape)
+    if power.ndim == 1:
+        power = power.view(*[1] * len(img_batch_dim), *power.shape)
+    if saturation.ndim == 0:
+        saturation = saturation.view(*[1] * len(img_batch_dim), *saturation.shape)
+
+    # Now check that the lengths are matching
+    assert slope.ndim == len(img_batch_dim) + 1
+    assert offset.ndim == len(img_batch_dim) + 1
+    assert power.ndim == len(img_batch_dim) + 1
+    assert saturation.ndim == len(img_batch_dim)
+
+    # Inverse Saturation adjustment
+    img = inverse_saturation(img, saturation)
+    # Inverse SOP adjustments
+    if clamp:
+        img = torch.clamp(img, 0.0, 1.0)
+    img = torch.where(
+        img > 1e-7, torch.pow(img, 1 / power.unsqueeze(-1).unsqueeze(-1)), img
+    )
+    img = (img - offset.unsqueeze(-1).unsqueeze(-1)) / slope.unsqueeze(-1).unsqueeze(-1)
+    # Clamp if specified
+    if clamp:
+        img = torch.clamp(img, 0.0, 1.0)
+    return img

src/utils/color_space.py

@@ -0,0 +1,80 @@
+import torch
+from jaxtyping import Float
+
+
+def srgb_to_linear(x: Float[torch.Tensor, "*B C"]) -> Float[torch.Tensor, "*B C"]:
+    switch_val = 0.04045
+    return torch.where(
+        torch.greater(x, switch_val),
+        ((x.clip(min=switch_val) + 0.055) / 1.055).pow(2.4),
+        x / 12.92,
+    )
+
+
+def linear_to_srgb(x: Float[torch.Tensor, "*B C"]) -> Float[torch.Tensor, "*B C"]:
+    switch_val = 0.0031308
+    return torch.where(
+        torch.greater(x, switch_val),
+        1.055 * x.clip(min=switch_val).pow(1.0 / 2.4) - 0.055,
+        x * 12.92,
+    )
+
+
+def rgb_to_lab(srgb: Float[torch.Tensor, "*B C"]) -> Float[torch.Tensor, "*B C"]:
+    srgb_pixels = torch.reshape(srgb, [-1, 3])
+
+    linear_mask = srgb_pixels <= 0.04045
+    exponential_mask = srgb_pixels > 0.04045
+    rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (
+        ((srgb_pixels + 0.055) / 1.055) ** 2.4
+    ) * exponential_mask
+
+    rgb_to_xyz = (
+        torch.tensor(
+            [
+                #    X        Y          Z
+                [0.412453, 0.212671, 0.019334],  # R
+                [0.357580, 0.715160, 0.119193],  # G
+                [0.180423, 0.072169, 0.950227],  # B
+            ]
+        )
+        .to(srgb.dtype)
+        .to(srgb.device)
+    )
+
+    xyz_pixels = torch.mm(rgb_pixels, rgb_to_xyz)
+
+    xyz_normalized_pixels = torch.mul(
+        xyz_pixels,
+        torch.tensor([1 / 0.950456, 1.0, 1 / 1.088754]).to(srgb.dtype).to(srgb.device),
+    )
+
+    epsilon = 6.0 / 29.0
+    linear_mask = (xyz_normalized_pixels <= (epsilon**3)).to(srgb.dtype).to(srgb.device)
+
+    exponential_mask = (
+        (xyz_normalized_pixels > (epsilon**3)).to(srgb.dtype).to(srgb.device)
+    )
+
+    fxfyfz_pixels = (
+        xyz_normalized_pixels / (3 * epsilon**2) + 4.0 / 29.0
+    ) * linear_mask + (
+        (xyz_normalized_pixels + 0.000001) ** (1.0 / 3.0)
+    ) * exponential_mask
+
+    fxfyfz_to_lab = (
+        torch.tensor(
+            [
+                #  l       a       b
+                [0.0, 500.0, 0.0],  # fx
+                [116.0, -500.0, 200.0],  # fy
+                [0.0, 0.0, -200.0],  # fz
+            ]
+        )
+        .to(srgb.dtype)
+        .to(srgb.device)
+    )
+    lab_pixels = torch.mm(fxfyfz_pixels, fxfyfz_to_lab) + torch.tensor(
+        [-16.0, 0.0, 0.0]
+    ).to(srgb.dtype).to(srgb.device)
+    return torch.reshape(lab_pixels, srgb.shape)

src/utils/image.py

@@ -0,0 +1,33 @@
+import cv2
+import numpy as np
+import torch
+from jaxtyping import Float
+
+
+def read_img(path):
+    img = cv2.imread(str(path), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
+    if img.ndim == 3:
+        img = cv2.cvtColor(img[..., :3], cv2.COLOR_BGR2RGB)
+    elif img.ndim == 2:
+        img = img[..., np.newaxis]
+    dinfo = np.iinfo(img.dtype)
+    return (img.astype(np.float32) / dinfo.max) * 2 - 1
+
+
+def write_img(path: str, data: np.ndarray):
+    data = np.clip(data * 0.5 + 0.5, 0, 1)
+    if data.ndim == 3 and data.shape[-1] == 3:
+        data = cv2.cvtColor(data, cv2.COLOR_RGB2BGR)
+    elif data.ndim == 2:
+        data = data[..., np.newaxis]
+
+    data = (data * 255).astype(np.uint8)
+    cv2.imwrite(path, data)
+
+
+def to_torch(img: Float[np.ndarray, "H W C"]) -> Float[torch.Tensor, "C H W"]:
+    return torch.from_numpy(img).permute(2, 0, 1)
+
+
+def from_torch(img: Float[torch.Tensor, "C H W"]) -> Float[np.ndarray, "H W C"]:
+    return img.permute(1, 2, 0).detach().cpu().float().numpy()

src/utils/math.py

@@ -0,0 +1,39 @@
+import math
+from typing import Optional, Union
+
+import torch
+
+
+def _random_orthonormal_matrix(d: int, device: torch.device) -> torch.Tensor:
+    """Draw a random rotation matrix Q ∈ SO(d) (Haar) via QR-factorisation."""
+    a = torch.randn(d, d, device=device)
+    # QR gives orthonormal columns; ensure right-handed
+    q, r = torch.linalg.qr(a, mode="reduced")
+    # make determinant +1 (special orthogonal) – flip first column if needed
+    if torch.det(q) < 0:
+        q[:, 0] = -q[:, 0]
+    return q  # (d,d)
+
+
+def sobol_sphere(
+    n: int,
+    d: int,
+    device: torch.device,
+    sobol_engine: Optional[torch.quasirandom.SobolEngine] = None,
+) -> Union[torch.Tensor, torch.quasirandom.SobolEngine]:
+    """n unit vectors on S^{d-1} via scrambled Sobol + Gaussian + random rotation."""
+    if sobol_engine is None:
+        sob = torch.quasirandom.SobolEngine(dimension=d, scramble=True)
+    else:
+        sob = sobol_engine
+    # Draw in [0,1)^d then map → 𝒩(0,1)
+    u01 = sob.draw(n).to(device)
+
+    eps = 1e-7
+    u01 = u01.clamp(min=eps, max=1.0 - eps)  # avoid 0 and 1 exactly
+
+    z = torch.erfinv(2.0 * u01 - 1.0) * math.sqrt(2.0)  # inverse-CDF of Normal
+    z = z / (z.norm(dim=1, keepdim=True) + 1e-8)  # project to sphere
+    # Random global rotation (RQMC) to make estimator unbiased
+    Q = _random_orthonormal_matrix(d, device)
+    return z @ Q.T, sob  # (n,d)