Initial commit: track binaries with LFS

.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
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,177 @@
+SD3_*
+*.pdf
+sd3_*
+wandb/*
+output/*
+data/*
+vscode/*
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+media/*
+pose/*
+wandb/*
+thirdparty/*
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+slurm*.out
+prompt_cache/*
+.vscode/
+cache/
+notebooks/
+paper_vis_scripts/
+sweeps/

DIV2k_mask.npy

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

LCFM/SD3.5-large_MSE_DIV2k.npy

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

README.md

@@ -1,14 +1,70 @@
----
-title: FLAIR
-emoji: 😻
-colorFrom: yellow
-colorTo: gray
-sdk: gradio
-sdk_version: 5.32.1
-app_file: app.py
-pinned: false
-license: mit
-short_description: Training-Free Super Resolution and Inpainting with FLAIR
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<div align="center">
+
+# FLAIR: Flow-Based Latent Alignment for Image Restoration
+
+**Julius Erbach<sup>1</sup>, Dominik Narnhofer<sup>1</sup>, Andreas Dombos<sup>1</sup>, Jan Eric Lenssen<sup>1</sup>, Bernt Schiele<sup>2</sup>, Konrad Schindler<sup>1</sup>**  
+<br>
+<sup>1</sup> Photogrammetry and Remote Sensing, ETH Zurich  
+<sup>2</sup> Max Planck Institute for Informatics, Saarbrücken  
+
+[![Paper](https://img.shields.io/badge/arXiv-PDF-b31b1b)](link)
+[![Page](https://img.shields.io/badge/Project-Page-green)](inverseFLAIR.github.io)
+[![Hugging Face Space](https://img.shields.io/badge/🤗%20Hugging%20Face-Space-yellow)](link)
+</div>
+
+<p align="center">
+  <img src="assets/teaser3.svg" alt="teaser" width=98%"/>
+</p>
+<p align="center">
+  <emph>FLAIR</emph> is a novel approach for solving inverse imaging problems using flow-based posterior sampling.
+</p>
+
+## Installation
+
+1.  Clone the repository:
+    ```bash
+    git clone <your-repo-url>
+    cd <your-repo-name>
+    ```
+
+2.  Create a virtual environment (recommended):
+    ```bash
+    python3 -m venv venv
+    source venv/bin/activate  
+    ```
+
+3.  Install the required dependencies from `requirements.txt`:
+    ```bash
+    pip install -r requirements.txt
+    pip install .
+    ```
+
+## Running Inference
+
+To run inference, you can use one of the Python script run_image_inv.py with the according config file.
+An example from the FFQH dataset
+```bash
+python inference_scripts/run_image_inv.py --config configs/inpainting.yaml --target_file examples/girl.png --result_folder output --prompt="a high quality photo of a face"
+```
+Or an example from the DIV2K dataset with captions provided by DAPE using the degraded input. The masks can be defined as rectanlge coordinates in the config file or provided as .npy file where true pixels are observed and false are masked out.
+
+```bash
+python inference_scripts/run_image_inv.py --config configs/inpainting.yaml --target_file examples/sunflowers.png --result_folder output --prompt="a high quality photo of bloom, blue, field, flower, sky, sunflower, sunflower field, yellow" --mask_file DIV2k_mask.npy
+```
+
+```bash
+python inference_scripts/run_image_inv.py --config configs/x12.yaml --target_file examples/sunflowers.png --result_folder output --prompt="a high quality photo of bloom, blue, field, flower, sky, sunflower, sunflower field, yellow"
+```
+
+## Citation
+
+If you find this work useful in your research, please cite our paper:
+
+```bibtex
+@article{er2025solving,
+  title={Solving Inverse Problems with FLAIR},
+  author={Erbach, Julius and Narnhofer, Dominik and Dombos, Andreas and Lenssen, Jan Eric and Schiele, Bernt and Schindler, Konrad},
+  journal={arXiv},
+  year={2025}
+}
+```

app.py

@@ -0,0 +1,774 @@
+import gradio as gr
+from gradio_imageslider import ImageSlider  # Replaces gr.ImageCompare
+import torch
+import yaml
+import numpy as np
+from PIL import Image
+import torchvision.transforms.functional as TF
+import random
+import os
+import sys
+import json # Added import
+import copy
+# Add project root to sys.path to allow direct import of var_post_samp
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "."))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+from flair.pipelines import model_loader
+from flair import var_post_samp, degradations
+
+
+
+CONFIG_FILE_PATH = "./configs/inpainting_gradio.yaml"
+DTYPE = torch.bfloat16
+
+# Global variables to hold the model and config
+MODEL = None
+POSTERIOR_MODEL = None
+BASE_CONFIG = None
+DEVICES = None
+PRIMARY_DEVICE = None
+# project_root is already defined globally, will be used by save_configuration
+
+SR_CONFIG_FILE_PATH = "./configs/x12_gradio.yaml"
+
+# Function to save the current configuration for demo examples
+def save_configuration(image_editor_data, image_input, prompt, seed_val, task, random_seed_bool, steps_val):
+    global project_root # Ensure access to the globally defined project_root
+    if task == "Super Resolution":
+        if image_input is None:
+            return gr.Markdown("""<p style='color:red;'>Error: No low-resolution image loaded.</p>""")
+        # For Super Resolution, we don't need a mask, just the image
+        input_image = image_input
+        mask_image = None
+    else:  # Inpainting task
+        if image_editor_data is None or image_editor_data['background'] is None:
+            return gr.Markdown("""<p style='color:red;'>Error: No background image loaded.</p>""")
+        
+        # Check if layers exist and the first layer (mask) is not None
+        if not image_editor_data['layers'] or image_editor_data['layers'][0] is None:
+            return gr.Markdown("""<p style='color:red;'>Error: No mask drawn. Please use the brush tool to draw a mask.</p>""")
+
+        input_image = image_editor_data['background']
+        mask_image = image_editor_data['layers'][0]
+
+    metadata = {
+        "prompt": prompt,
+        "seed_on_slider": int(seed_val),
+        "use_random_seed_checkbox": bool(random_seed_bool),
+        "num_steps": int(steps_val),
+        "task_type": task  # Always inpainting for now
+    }
+
+    demo_images_dir = os.path.join(project_root, "demo_images")
+    try:
+        os.makedirs(demo_images_dir, exist_ok=True)
+    except Exception as e:
+        return gr.Markdown(f"""<p style='color:red;'>Error creating directory {demo_images_dir}: {str(e)}</p>""")
+
+    i = 0
+    while True:
+        base_filename = f"demo_{i}"
+        meta_check_path = os.path.join(demo_images_dir, f"{base_filename}_meta.json")
+        if not os.path.exists(meta_check_path):
+            break
+        i += 1
+    
+    image_save_path = os.path.join(demo_images_dir, f"{base_filename}_image.png")
+    mask_save_path = os.path.join(demo_images_dir, f"{base_filename}_mask.png")
+    meta_save_path = os.path.join(demo_images_dir, f"{base_filename}_meta.json")
+
+    try:
+        input_image.save(image_save_path)
+        if mask_image is not None:
+            # Ensure mask is saved in a usable format, e.g., 'L' mode for grayscale, or 'RGBA' if it has transparency
+            if mask_image.mode != 'L' and mask_image.mode != '1': # If not already grayscale or binary
+                mask_image = mask_image.convert('RGBA') # Preserve transparency if drawn, or convert to L
+            mask_image.save(mask_save_path)
+        
+        with open(meta_save_path, 'w') as f:
+            json.dump(metadata, f, indent=4)
+        return gr.Markdown(f"""<p style='color:green;'>Configuration saved as {base_filename} in demo_images folder.</p>""")
+    except Exception as e:
+        return gr.Markdown(f"""<p style='color:red;'>Error saving configuration: {str(e)}</p>""")
+
+def embed_prompt(prompt, device):
+    print(f"Generating prompt embeddings for: {prompt}")
+    with torch.no_grad(): # Add torch.no_grad() here
+        POSTERIOR_MODEL.model.text_encoder.to(device).to(torch.bfloat16)
+        POSTERIOR_MODEL.model.text_encoder_2.to(device).to(torch.bfloat16)
+        POSTERIOR_MODEL.model.text_encoder_3.to(device).to(torch.bfloat16)
+        (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+        ) = POSTERIOR_MODEL.model.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt,
+            prompt_3=prompt,
+            negative_prompt="",
+            negative_prompt_2="",
+            negative_prompt_3="",
+            do_classifier_free_guidance=POSTERIOR_MODEL.model.do_classifier_free_guidance,
+            prompt_embeds=None,
+            negative_prompt_embeds=None,
+            pooled_prompt_embeds=None,
+            negative_pooled_prompt_embeds=None,
+            device=device,
+            clip_skip=None,
+            num_images_per_prompt=1,
+            max_sequence_length=256,
+            lora_scale=None,
+    )
+    # POSTERIOR_MODEL.model.text_encoder.to("cpu").to(torch.bfloat16)
+    # POSTERIOR_MODEL.model.text_encoder_2.to("cpu").to(torch.bfloat16)
+    # POSTERIOR_MODEL.model.text_encoder_3.to("cpu").to(torch.bfloat16)
+    torch.cuda.empty_cache()  # Clear GPU memory after embedding generation
+    return {
+        "prompt_embeds": prompt_embeds.to(device, dtype=DTYPE),
+        "negative_prompt_embeds": negative_prompt_embeds.to(device, dtype=DTYPE) if negative_prompt_embeds is not None else None,
+        "pooled_prompt_embeds": pooled_prompt_embeds.to(device, dtype=DTYPE),
+        "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds.to(device, dtype=DTYPE) if negative_pooled_prompt_embeds is not None else None
+    }
+
+def initialize_globals():
+    global MODEL, POSTERIOR_MODEL, BASE_CONFIG, DEVICES, PRIMARY_DEVICE
+
+    print("Global initialization started...")
+    # Setup device (run once)
+    if torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+        DEVICES = [f"cuda:{i}" for i in range(num_gpus)]
+        PRIMARY_DEVICE = DEVICES[0]
+        print(f"Initializing with devices: {DEVICES}, Primary: {PRIMARY_DEVICE}")
+    else:
+        DEVICES = ["cpu"]
+        PRIMARY_DEVICE = "cpu"
+        print("No CUDA devices found. Initializing with CPU.")
+
+    # Load base configuration (once)
+    with open(CONFIG_FILE_PATH, "r") as f:
+        BASE_CONFIG = yaml.safe_load(f)
+    
+    # Prepare a temporary config for the initial model and posterior_model loading
+    init_config = BASE_CONFIG.copy()
+    
+    # Ensure prompt/caption settings are valid for model_loader for initialization
+    # Forcing prompt mode for initial load.
+    init_config["prompt"] = [BASE_CONFIG.get("prompt", "Initialization prompt")]
+    init_config["caption_file"] = None
+    
+    # Default values that might be needed by model_loader or utils called within
+    init_config.setdefault("target_file", "dummy_target.png") 
+    init_config.setdefault("result_file", "dummy_results/")
+    init_config.setdefault("seed", random.randint(0, 2**32 - 1)) # Init with a random seed
+
+    print("Loading base model and variational posterior model once...")
+    # MODEL is the main diffusion model, loaded once.
+    # inp_kwargs_for_init are based on init_config, not directly used for subsequent inferences.
+    model_obj, _ = model_loader.load_model(init_config, device=DEVICES)
+    MODEL = model_obj
+
+    # Initialize VariationalPosterior once with the loaded MODEL and init_config.
+    # Its internal forward_operator will be based on init_config's degradation settings,
+    # but will be replaced in each inpaint_image call.
+    POSTERIOR_MODEL = var_post_samp.VariationalPosterior(MODEL, init_config)
+    print("Global initialization complete.")
+
+
+def load_config_for_inference(prompt_text, seed=None):
+    # This function is now for creating a temporary config for each inference call,
+    # primarily to get up-to-date inp_kwargs via model_loader.
+    # It starts from BASE_CONFIG and applies current overrides.
+    if BASE_CONFIG is None:
+        raise RuntimeError("Base config not initialized. Call initialize_globals().")
+
+    current_config = BASE_CONFIG.copy()
+
+    current_config["prompt"] = [prompt_text] # Override with user's prompt
+    current_config["caption_file"] = None # Ensure we are in prompt mode
+
+    if seed is None:
+        seed = current_config.get("seed", random.randint(0, 2**32 - 1))
+    current_config["seed"] = seed
+    # Set global seeds for reproducibility for the current call
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    print(f"Using seed for current inference: {seed}")
+    
+    # Ensure other necessary fields are in 'current_config' if model_loader needs them
+    current_config.setdefault("target_file", "dummy_target.png") 
+    current_config.setdefault("result_file", "dummy_results/")
+
+    return current_config
+
+def preprocess_image(pil_image, resolution, is_mask=False):
+    img = pil_image.convert("RGB") if not is_mask else pil_image.convert("L")
+
+    # Calculate new dimensions to maintain aspect ratio, making shorter edge 'resolution'
+    original_width, original_height = img.size
+    if original_width < original_height:
+        new_short_edge = resolution
+        new_long_edge = int(resolution * (original_height / original_width))
+        new_width = new_short_edge
+        new_height = new_long_edge
+    else:
+        new_short_edge = resolution
+        new_long_edge = int(resolution * (original_width / original_height))
+        new_height = new_short_edge
+        new_width = new_long_edge
+
+    # TF.resize expects [height, width]
+    img = TF.resize(img, [new_height, new_width], interpolation=TF.InterpolationMode.LANCZOS)
+
+    # Center crop to the target square resolution
+    img = TF.center_crop(img, [resolution, resolution])
+
+    img_tensor = TF.to_tensor(img) # Scales to [0, 1]
+    if is_mask:
+        # Ensure mask is binary (0 or 1), 1 for region to inpaint
+        # The mask from ImageEditor is RGBA, convert to L first.
+        img = img.convert('L')
+        img_tensor = TF.to_tensor(img) # Recalculate tensor after convert
+        img_tensor = (img_tensor == 0.) # Threshold for mask (drawn parts are usually non-black)
+        img_tensor = img_tensor.repeat(3, 1, 1) # Repeat mask across 3 channels
+    else:
+        # Normalize image to [-1, 1]
+        img_tensor = img_tensor * 2 - 1 
+    return img_tensor.unsqueeze(0) # Add batch dimension
+
+def preprocess_lr_image(pil_image, resolution, device, dtype):
+    if pil_image is None:
+        raise ValueError("Input PIL image cannot be None.")
+    img = pil_image.convert("RGB")
+
+    # Center crop to the target square resolution (no resizing)
+    img = TF.center_crop(img, [resolution, resolution])
+
+    img_tensor = TF.to_tensor(img)  # Scales to [0, 1]
+    # Normalize image to [-1, 1]
+    img_tensor = img_tensor * 2 - 1 
+    return img_tensor.unsqueeze(0).to(device, dtype=dtype) # Add batch dimension and move to device
+
+
+def postprocess_image(image_tensor):
+    # Remove batch dimension, move to CPU, convert to float
+    image_tensor = image_tensor.squeeze(0).cpu().float()
+    # Denormalize from [-1, 1] to [0, 1]
+    image_tensor = image_tensor * 0.5 + 0.5
+    # Clip values to [0, 1]
+    image_tensor = torch.clamp(image_tensor, 0, 1)
+    # Convert to PIL Image
+    pil_image = TF.to_pil_image(image_tensor)
+    return pil_image
+
+def inpaint_image(image_editor_output, prompt_text, fixed_seed_value, use_random_seed, guidance_scale, num_steps): # MODIFIED: seed_input changed to fixed_seed_value, use_random_seed
+    try:
+        if image_editor_output is None:
+            raise gr.Error("Please upload an image and draw a mask.")
+
+        input_pil = image_editor_output['background'] 
+        
+        if not image_editor_output['layers'] or image_editor_output['layers'][0] is None:
+            raise gr.Error("Please draw a mask on the image using the brush tool.")
+        mask_pil = image_editor_output['layers'][0]
+
+
+        if input_pil is None:
+            raise gr.Error("Please upload an image.")
+        if mask_pil is None: 
+            raise gr.Error("Please draw a mask on the image.")
+
+        current_seed = None
+        if use_random_seed:
+            current_seed = None # load_config_for_inference will generate a random seed
+        else:
+            try:
+                current_seed = int(fixed_seed_value)
+            except ValueError:
+                # This should ideally not happen with a slider, but good for robustness
+                raise gr.Error("Seed must be an integer.")
+
+        # Prepare config for current inference (gets prompt, seed)
+        current_config = load_config_for_inference(prompt_text, current_seed)
+        resolution = current_config["resolution"]
+
+        # MODIFIED: Set num_steps from slider into the current_config
+        # Assuming 'num_steps' is a key POSTERIOR_MODEL will use from its config.
+        # Common alternatives could be current_config['solver_kwargs']['n_steps'] = num_steps
+        current_config['n_steps'] = int(num_steps) 
+        print(f"Using num_steps: {current_config['n_steps']}")
+
+
+        # Preprocess image and mask
+        guidance_img_tensor = preprocess_image(input_pil, resolution, is_mask=False).to(PRIMARY_DEVICE, dtype=DTYPE)
+        # Mask from ImageEditor is RGBA, preprocess_image will handle conversion to L and then binary
+        mask_tensor = preprocess_image(mask_pil, resolution, is_mask=True).to(PRIMARY_DEVICE, dtype=DTYPE) 
+        
+        # Get inp_kwargs for the CURRENT prompt and config.
+        print("Preparing inference inputs (e.g., prompt embeddings)...")
+        prompt_embeds = embed_prompt(prompt_text, device=PRIMARY_DEVICE) # Embed the prompt for the current inference
+        current_inp_kwargs = prompt_embeds
+        # MODIFIED: Use guidance_scale from slider
+        current_inp_kwargs['guidance'] = float(guidance_scale) 
+        print(f"Using guidance_scale: {current_inp_kwargs['guidance']}")
+        
+        # Update the global POSTERIOR_MODEL's config for this call.
+        # This ensures its methods use the latest settings (like num_steps) if they access self.config.
+        POSTERIOR_MODEL.config = current_config
+        POSTERIOR_MODEL.model._guidance_scale = guidance_scale
+        print("Applying forward operator (masking)...")
+        # Directly set the forward_operator on the global POSTERIOR_MODEL instance
+        # H and W are height and width of the guidance image tensor
+        POSTERIOR_MODEL.forward_operator = degradations.Inpainting(
+            mask=mask_tensor.bool()[0], # Inpainting often expects a boolean mask
+            H=guidance_img_tensor.shape[2], 
+            W=guidance_img_tensor.shape[3], 
+            noise_std=0, 
+        )
+        y = POSTERIOR_MODEL.forward_operator(guidance_img_tensor)    
+        
+        print("Running inference...")
+        with torch.no_grad():
+            # Use the global POSTERIOR_MODEL instance
+            result_dict = POSTERIOR_MODEL.forward(y, current_inp_kwargs)
+        
+        x_hat = result_dict["x_hat"]
+        
+        print("Postprocessing result...")
+        output_pil = postprocess_image(x_hat)
+        
+        # Convert mask tensor to PIL image for display
+        # Mask tensor is [0, 1], take one channel, convert to PIL
+        mask_display_tensor = mask_tensor.squeeze(0).cpu().float() # Remove batch, move to CPU
+        # If mask_tensor was (B, 3, H, W) and binary 0 or 1 (after repeat)
+        # We can take any channel, e.g., mask_display_tensor[0]
+        # Ensure it's (H, W) or (1, H, W) for to_pil_image
+        if mask_display_tensor.ndim == 3 and mask_display_tensor.shape[0] == 3: # (C, H, W)
+            mask_display_tensor = mask_display_tensor[0] # Take one channel (H, W)
+        
+        # Ensure it's in the range [0, 1] and suitable for PIL conversion
+        # If it was 0. for masked and 1. for unmasked (or vice-versa depending on logic)
+        # TF.to_pil_image expects [0,1] for single channel float
+        mask_pil_display = TF.to_pil_image(mask_display_tensor)
+
+        return output_pil, [output_pil, output_pil], current_config["seed"] # MODIFIED: Removed mask_pil_display
+    except gr.Error as e: # Handle Gradio-specific errors first
+        raise
+    except Exception as e:
+        print(f"Error during inpainting: {e}")
+        import traceback # Ensure traceback is imported here if not globally
+        traceback.print_exc()
+        # Return a more user-friendly error message to Gradio
+        raise gr.Error(f"An error occurred: {str(e)}. Check console for details.")
+
+def super_resolution_image(lr_image, prompt_text, fixed_seed_value, use_random_seed, guidance_scale, num_steps, sr_scale_factor, downscale_input):
+    try:
+        if lr_image is None:
+            raise gr.Error("Please upload a low-resolution image.")
+
+        current_seed = None
+        if use_random_seed:
+            current_seed = random.randint(0, 2**32 - 1)
+        else:
+            try:
+                current_seed = int(fixed_seed_value)
+            except ValueError:
+                raise gr.Error("Seed must be an integer.")
+
+        # Load Super-Resolution specific configuration
+        if not os.path.exists(SR_CONFIG_FILE_PATH):
+            raise gr.Error(f"Super-resolution config file not found: {SR_CONFIG_FILE_PATH}")
+        with open(SR_CONFIG_FILE_PATH, "r") as f:
+            sr_base_config = yaml.safe_load(f)
+
+        current_sr_config = copy.deepcopy(sr_base_config) # Start with a copy of the base SR config
+        current_sr_config["prompt"] = [prompt_text]
+        current_sr_config["caption_file"] = None # Ensure prompt mode
+        current_sr_config["seed"] = current_seed
+        
+        torch.manual_seed(current_seed)
+        np.random.seed(current_seed)
+        random.seed(current_seed)
+        print(f"Using seed for SR inference: {current_seed}")
+
+        current_sr_config['n_steps'] = int(num_steps)
+        current_sr_config["degradation"]["kwargs"]["scale"] = sr_scale_factor
+        current_sr_config["optimizer_dataterm"]["kwargs"]["lr"] = sr_base_config.get("optimizer_dataterm", {}).get("kwargs", {}).get("lr") * sr_scale_factor**2 / (sr_base_config.get("degradation", {}).get("kwargs", {}).get("scale")**2)
+        print(f"Using num_steps for SR: {current_sr_config['n_steps']}")
+
+        # Determine target HR resolution for the output
+        hr_resolution = current_sr_config.get("degradation", {}).get("kwargs", {}).get("img_size")
+        # Calculate target LR dimensions based on the chosen scale factor
+        target_lr_width = int(hr_resolution / sr_scale_factor)
+        target_lr_height = int(hr_resolution / sr_scale_factor)
+        print(f"Target LR dimensions for SR: {target_lr_width}x{target_lr_height} for scale x{sr_scale_factor}")
+
+        print("Preparing SR inference inputs (prompt embeddings)...")
+        prompt_embeds = embed_prompt(prompt_text, device=PRIMARY_DEVICE)
+        current_inp_kwargs = prompt_embeds
+        current_inp_kwargs['guidance'] = float(guidance_scale)
+        print(f"Using guidance_scale for SR: {current_inp_kwargs['guidance']}")
+
+        POSTERIOR_MODEL.config = current_sr_config
+        POSTERIOR_MODEL.model._guidance_scale = float(guidance_scale)
+
+        print("Applying SR forward operator...")
+        
+        POSTERIOR_MODEL.forward_operator = degradations.SuperResGradio(
+            **current_sr_config["degradation"]["kwargs"]
+        )
+        
+        if downscale_input:
+            y_tensor = preprocess_lr_image(lr_image, hr_resolution, PRIMARY_DEVICE, DTYPE)
+            # y_tensor = POSTERIOR_MODEL.forward_operator(y_tensor)
+            y_tensor = torch.nn.functional.interpolate(y_tensor, scale_factor=1/sr_scale_factor, mode='bilinear', align_corners=False, antialias=True)
+            # simulate 8bit input by quantizing to 8-bit
+            y_tensor = ((y_tensor * 127.5 + 127.5).clamp(0, 255).to(torch.uint8) / 127.5 - 1.0).to(DTYPE) 
+        else:
+            # check if the input image has the correct dimensions
+            if lr_image.size[0] != target_lr_width or lr_image.size[1] != target_lr_height:
+                raise gr.Error(f"Input image must be {target_lr_width}x{target_lr_height} pixels for the selected scale factor of {sr_scale_factor}.")
+            y_tensor = preprocess_lr_image(lr_image, target_lr_width, PRIMARY_DEVICE, DTYPE)
+            # add some noise to the input image
+            noise_std = current_sr_config.get("degradation", {}).get("kwargs", {}).get("noise_std", 0.0)
+            y_tensor += torch.randn_like(y_tensor) * noise_std
+            # save for debugging purposes
+            # first convert to PIL
+            pil_y = postprocess_image(y_tensor)# Remove batch dimension and convert to PIL
+            pil_y.save("debug_input_image.png")  # Save the input image for debugging
+                
+        
+        print("Running SR inference...")
+        with torch.no_grad():
+            result_dict = POSTERIOR_MODEL.forward(y_tensor, current_inp_kwargs)
+        
+        x_hat = result_dict["x_hat"]
+        
+        print("Postprocessing SR result...")
+        output_pil = postprocess_image(x_hat)
+
+        # Upscale input image with nearest neighbor for comparison
+        upscaled_input = y_tensor.reshape(1,3,target_lr_height, target_lr_width)
+        upscaled_input = POSTERIOR_MODEL.forward_operator.nn(upscaled_input)  # Use nearest neighbor upscaling
+        upscaled_input = postprocess_image(upscaled_input)
+        # save for debugging purposes
+        upscaled_input.save("debug_upscaled_input.png")  # Save the upscaled input image for debugging
+        # upscaled_input = upscaled_input.resize((hr_resolution, hr_resolution), resample=Image.NEAREST)
+        return (upscaled_input, output_pil), current_sr_config["seed"]
+
+    except gr.Error as e:
+        raise
+    except Exception as e:
+        print(f"Error during super-resolution: {e}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"An error occurred during super-resolution: {str(e)}. Check console for details.")
+
+
+# Input for seed, allowing users to set it or leave it blank for random/config default
+# Determine default num_steps from BASE_CONFIG if available
+default_num_steps = 50 # Fallback default
+if BASE_CONFIG is not None: # Check if BASE_CONFIG has been initialized
+    default_num_steps = BASE_CONFIG.get("num_steps", BASE_CONFIG.get("solver_kwargs", {}).get("num_steps", 50))
+
+def superres_preview_preprocess(pil_image, resolution=768):
+    if pil_image is None:
+        return None
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+    # check if image is smaller than resolution
+    original_width, original_height = pil_image.size
+    if original_width < resolution or original_height < resolution:
+        return pil_image  # No resizing needed, return original image
+    else:
+        pil_image = TF.center_crop(pil_image, [resolution, resolution])
+    return pil_image
+
+
+# Dynamically load examples from demo_images directory
+example_list_inp = []
+example_list_sr = []
+demo_images_dir = os.path.join(project_root, "static/demo_images")
+
+if os.path.exists(demo_images_dir):
+    filenames = sorted(os.listdir(demo_images_dir))
+    processed_bases = set()
+    for filename in filenames:
+        if filename.startswith("demo_") and filename.endswith("_meta.json"):
+            base_name = filename[:-len("_meta.json")] # e.g., "demo_0"
+            if base_name in processed_bases:
+                continue
+
+            meta_path = os.path.join(demo_images_dir, filename)
+            image_filename = f"{base_name}_image.png"
+            image_path = os.path.join(demo_images_dir, image_filename)
+            mask_filename = f"{base_name}_mask.png"
+            mask_path = os.path.join(demo_images_dir, mask_filename)
+
+            if os.path.exists(image_path):
+                try:
+                    with open(meta_path, 'r') as f:
+                        metadata = json.load(f)
+                    task = metadata.get("task_type")
+                    prompt = metadata.get("prompt", "")
+                    if task == "Super Resolution":
+                        example_list_sr.append([image_path, prompt, task])
+                    else:
+                        image_editor_input = {
+                            "background": image_path,
+                            "layers": [mask_path],
+                            "composite": None  # Add this key to satisfy ImageEditor's as_example processing
+                        }
+                        example_list_inp.append([image_editor_input, prompt, task])
+                    
+                    # Structure for ImageEditor: { "background": filepath, "layers": [filepath], "composite": None }
+                    
+                except json.JSONDecodeError:
+                    print(f"Warning: Could not decode JSON from {meta_path}. Skipping example {base_name}.")
+                except Exception as e:
+                    print(f"Warning: Error processing example {base_name}: {e}. Skipping.")
+            else:
+                missing_files = []
+                if not os.path.exists(image_path):
+                    missing_files.append(image_filename)
+                if not os.path.exists(mask_path):
+                    missing_files.append(mask_filename)
+                print(f"Warning: Missing files for example {base_name} ({', '.join(missing_files)}). Skipping.")
+else:
+    print(f"Info: 'demo_images' directory not found at {demo_images_dir}. No dynamic examples will be loaded.")
+
+
+if __name__ == "__main__":
+    if not os.path.exists(CONFIG_FILE_PATH):
+        print(f"ERROR: Configuration file not found at {CONFIG_FILE_PATH}")
+        sys.exit(1)
+    
+    initialize_globals()
+    
+    if MODEL is None or POSTERIOR_MODEL is None:
+        print("ERROR: Global model initialization failed.")
+        sys.exit(1)
+
+    # --- Define Gradio UI using gr.Blocks after globals are initialized ---
+    title_str = "Solving Inverse Problems with FLAIR: Inpainting Demo"
+    description_str = """
+Select a task (Inpainting or Super Resolution) and upload an image.
+For Inpainting, draw a mask on the image to specify the area to be filled. We observed that our model can event solve simple editing task, if provided with an appropriate prompt. For large masks the step size might need to be adjusted to e.g. 80.
+For Super Resolution, upload a low-resolution image and select the upscaling factor. Images are always upscaled to 768x768 pixels. Therefore, for x12 superresolution, the input image must be 64x64 pixels. You can also upload a high resolution image which will be downscaled to the correct input size. 
+Use the slider to compare the low resolution input image with the super-resolved output.
+
+"""
+
+    # Determine default values now that BASE_CONFIG is initialized
+    default_num_steps = BASE_CONFIG.get("num_steps", BASE_CONFIG.get("solver_kwargs", {}).get("num_steps", 50))
+    default_guidance_scale = BASE_CONFIG.get("guidance", 2.0)
+
+    with gr.Blocks() as iface:
+        gr.Markdown(f"## {title_str}")
+        gr.Markdown(description_str)
+
+        task_selector = gr.Dropdown(
+            choices=["Inpainting", "Super Resolution"],
+            value="Inpainting",
+            label="Task"
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):  # Input column
+                # Inpainting Inputs
+                image_editor = gr.ImageEditor(
+                    type="pil",
+                    label="Upload Image & Draw Mask (for Inpainting)",
+                    sources=["upload"],
+                    height=512,
+                    width=512,
+                    visible=True
+                )
+
+                # Super Resolution Inputs
+                image_input = gr.Image(
+                    type="pil",
+                    label="Upload Low-Resolution Image (for Super Resolution)",
+                    visible=False
+                )
+
+                sr_scale_slider = gr.Dropdown(
+                    choices=[2, 4, 8, 12, 24],
+                    value=12,
+                    label="Upscaling Factor (Super Resolution)",
+                    interactive=True,
+                    visible=False # Initially hidden
+                )
+                downscale_input = gr.Checkbox(
+                    label="Downscale the provided image.",
+                    value=True,
+                    interactive=True,
+                    visible=False # Initially hidden
+                )
+
+                # Common Inputs
+                prompt_text = gr.Textbox(
+                    label="Prompt",
+                    placeholder="E.g., a beautiful landscape, a detailed portrait"
+                )
+                seed_slider = gr.Slider(
+                    minimum=0,
+                    maximum=2**32 -1, # Max for torch.manual_seed
+                    step=1,
+                    label="Seed (if not random)",
+                    value=42,
+                    interactive=True
+                )
+                
+                use_random_seed_checkbox = gr.Checkbox(
+                    label="Use Random Seed",
+                    value=True,
+                    interactive=True
+                )
+                guidance_scale_slider = gr.Slider(
+                    minimum=1.0,
+                    maximum=15.0,
+                    step=0.5,
+                    value=default_guidance_scale,
+                    label="Guidance Scale"
+                )
+                num_steps_slider = gr.Slider(
+                    minimum=28,
+                    maximum=150,
+                    step=1,
+                    value=default_num_steps,
+                    label="Number of Steps"
+                )
+                submit_button = gr.Button("Submit")
+
+                # # Add Save Configuration button and status text
+                # gr.Markdown("---") # Separator
+                # save_button = gr.Button("Save Current Configuration for Demo")
+                # save_status_text = gr.Markdown()
+
+            with gr.Column(scale=1):  # Output column
+                output_image_display = gr.Image(type="pil", label="Result")
+                sr_compare_display = ImageSlider(label="Super-Resolution: Input vs Output", visible=False, position=0.5)
+
+
+        
+
+        # --- Task routing and visibility logic ---
+        def update_visibility(task):
+            is_inpainting = task == "Inpainting"
+            is_super_resolution = task == "Super Resolution"
+            return {
+                image_editor: gr.update(visible=is_inpainting),
+                image_input: gr.update(visible=is_super_resolution),
+                sr_scale_slider: gr.update(visible=is_super_resolution),
+                downscale_input: gr.update(visible=is_super_resolution),
+                output_image_display: gr.update(visible=is_inpainting),
+                sr_compare_display: gr.update(visible=is_super_resolution, position=0.5),
+                downscale_input: gr.update(visible=is_super_resolution),
+            }
+
+        task_selector.change(
+            fn=update_visibility,
+            inputs=[task_selector],
+            outputs=[image_editor, image_input, sr_scale_slider, downscale_input, output_image_display, sr_compare_display]
+        )
+
+
+        # MODIFIED route_task to accept sr_scale_factor
+        def route_task(task, image_editor_data, lr_image_for_sr, prompt_text, fixed_seed_value, use_random_seed, guidance_scale, num_steps, sr_scale_factor_value, downscale_input):
+            if task == "Inpainting":
+                return inpaint_image(image_editor_data, prompt_text, fixed_seed_value, use_random_seed, guidance_scale, num_steps)
+            elif task == "Super Resolution":
+                result_images, seed_val = super_resolution_image(
+                    lr_image_for_sr, prompt_text, fixed_seed_value, use_random_seed,
+                    guidance_scale, num_steps, sr_scale_factor_value, downscale_input
+                )
+                return result_images[1], gr.update(value=result_images, position=0.5), seed_val
+            else:
+                raise gr.Error("Unsupported task.")
+
+        submit_button.click(
+            fn=route_task,
+            inputs=[
+                task_selector,
+                image_editor,
+                image_input,
+                prompt_text,
+                seed_slider,
+                use_random_seed_checkbox,
+                guidance_scale_slider,
+                num_steps_slider,
+                sr_scale_slider,
+                downscale_input,
+            ],
+            outputs=[
+                output_image_display,
+                sr_compare_display,
+                seed_slider
+            ]
+        )
+
+        # Wire up the save button
+        # save_button.click(
+        #     fn=save_configuration,
+        #     inputs=[
+        #         image_editor,
+        #         image_input,
+        #         prompt_text,
+        #         seed_slider,
+        #         task_selector,
+        #         use_random_seed_checkbox,
+        #         num_steps_slider,
+        #     ],
+        #     outputs=[save_status_text]
+        # )
+
+
+        gr.Markdown("---") # Separator
+        gr.Markdown("### Click an example to load:")
+        def load_example(example_data, prompt, task):
+            image_editor_input = example_data[0]
+            prompt_value = example_data[1]
+            if task == "Inpainting":
+                image_editor.clear()  # Clear current image and mask
+                if image_editor_input and image_editor_input.get("background"):
+                    image_editor.upload_image(image_editor_input["background"])
+                if image_editor_input and image_editor_input.get("layers"):
+                    for layer in image_editor_input["layers"]:
+                        image_editor.upload_mask(layer)
+            elif task == "Super Resolution":
+                image_input.clear()
+                image_input.upload_image(image_editor_input)
+                
+            # Set the prompt
+            prompt_text.value = prompt_value
+            # Optionally, set a random seed and guidance scale
+            seed_slider.value = random.randint(0, 2**32 - 1)
+            guidance_scale_slider.value = default_guidance_scale
+            # Set the task selector from the example
+            task_selector.set_value(task)
+            update_visibility(task)  # Update visibility based on task
+
+        with gr.Row():
+            gr.Examples(
+                examples=example_list_sr,
+                inputs=[image_input, prompt_text, task_selector],
+                label="Super Resolution Examples",
+                fn=load_example,
+            )
+        with gr.Row():
+            gr.Examples(
+                examples=example_list_inp,
+                inputs=[image_editor, prompt_text, task_selector],
+                label="Inpainting Examples",
+                fn=load_example,
+            )
+
+    # --- End of Gradio UI definition ---
+
+    print("Launching Gradio demo...")
+    iface.launch()

configs/inpainting.yaml

@@ -0,0 +1,45 @@
+# === Model & Data Settings ===
+model: "SD3"
+resolution: 768
+lambda_func: v
+optimized_reg_weight: ./LCFM/SD3_loss_v_MSE_DIV2k_neg_prompt.npy
+regularizer_weight: 0.5
+likelihood_weight_mode: reg_weight
+likelihood_steps: 15
+early_stopping: 1.e-4
+epochs: 1
+guidance: 2
+quantize: False
+use_tiny_ae: True
+negative_prompt: ""
+reg-shift: 0.0  
+
+# === Optimization Settings ===
+optimizer: 
+  name: SGD
+  kwargs:
+    lr: 1
+optimizer_dataterm: 
+  name: SGD
+  kwargs:
+    lr: 0.1
+
+# === Sampling & Misc ===
+t_sampling: descending
+n_steps: 50
+inv_alpha: 1-t
+ts_min: 0.18
+projection: False
+seed: 42
+
+# === Experiment-Specific Settings ===
+prompt: A high quality photo of a girl with a pirate eye-patch.
+degradation:
+  name: Inpainting
+  kwargs:
+    mask: [128, 640, 384, 640]
+    H: 768
+    W: 768
+    noise_std: 0.01
+likelihood_weight: 1
+

configs/inpainting_gradio.yaml

@@ -0,0 +1,45 @@
+# === Model & Data Settings ===
+model: "SD3"
+resolution: 512
+lambda_func: v
+optimized_reg_weight: ./LCFM/SD3_loss_v_MSE_DIV2k_neg_prompt.npy
+regularizer_weight: 0.5
+likelihood_weight_mode: reg_weight
+likelihood_steps: 15
+early_stopping: 1.e-4
+epochs: 1
+guidance: 2
+quantize: False
+use_tiny_ae: True
+negative_prompt: ""
+reg-shift: 0.0  
+
+# === Optimization Settings ===
+optimizer: 
+  name: SGD
+  kwargs:
+    lr: 1
+optimizer_dataterm: 
+  name: SGD
+  kwargs:
+    lr: 0.1
+
+# === Sampling & Misc ===
+t_sampling: descending
+n_steps: 50
+inv_alpha: 1-t
+ts_min: 0.18
+projection: False
+seed: 42
+
+# === Experiment-Specific Settings ===
+prompt: A high quality photo of a girl with a pirate eye-patch.
+degradation:
+  name: Inpainting
+  kwargs:
+    mask: [128, 640, 384, 640]
+    H: 512
+    W: 512
+    noise_std: 0.0
+likelihood_weight: 1
+

configs/motion_deblur.yaml

@@ -0,0 +1,43 @@
+# === Model & Data Settings ===
+model: "SD3"
+resolution: 768
+lambda_func: v
+optimized_reg_weight: SD3_loss_v_norm_div2k_squared.npy
+regularizer_weight: 0.5
+likelihood_weight_mode: reg_weight
+likelihood_steps: 15
+early_stopping: 1.e-4
+epochs: 1
+guidance: 2
+quantize: False
+use_tiny_ae: True
+negative_prompt: ""
+reg-shift: 0.0
+
+# === Optimization Settings ===
+optimizer: 
+  name: SGD
+  kwargs:
+    lr: 1 
+optimizer_dataterm: 
+  name: SGD
+  kwargs:
+    lr: 0.1
+
+# === Sampling & Misc ===
+t_sampling: descending
+n_steps: 50
+inv_alpha: 1-t
+ts_min: 0.18
+projection: False
+seed: 42
+
+# === Experiment-Specific Settings ===
+prompt: a high quality photo of a face
+degradation: 
+  name: MotionBlur
+  kwargs:
+    kernel_size: 61
+    noise_std: 0.01
+    img_size: 768
+likelihood_weight: 1

configs/x12.yaml

@@ -0,0 +1,48 @@
+# OmegaConf base config for neurips experiments
+# Place common settings here. Each experiment config will override as needed.
+
+# === Model & Data Settings ===
+model: "SD3"
+resolution: 768
+lambda_func: v
+optimized_reg_weight: ./LCFM/SD3_loss_v_MSE_DIV2k_neg_prompt.npy
+regularizer_weight: 0.5
+likelihood_weight_mode: reg_weight
+likelihood_steps: 15
+early_stopping: 1.e-4
+epochs: 1
+guidance: 2
+quantize: False
+use_tiny_ae: True
+negative_prompt: ""
+reg-shift: 0.0  
+
+# === Optimization Settings ===
+optimizer: 
+  name: SGD
+  kwargs:
+    lr: 1
+optimizer_dataterm: 
+  name: SGD
+  kwargs:
+    lr: 12
+
+# === Sampling & Misc ===
+t_sampling: descending
+n_steps: 50
+inv_alpha: 1-t
+quantize: False
+ts_min: 0.18
+projection: False
+seed: 3
+
+# === Experiment-Specific Settings ===
+prompt: a high quality photo of
+# Super-resolution x12 for FFHQ
+degradation:
+  name: SuperRes
+  kwargs:
+    scale: 12
+    noise_std: 0.01
+    img_size: 768
+likelihood_weight: 1

configs/x12_gradio.yaml

@@ -0,0 +1,48 @@
+# OmegaConf base config for neurips experiments
+# Place common settings here. Each experiment config will override as needed.
+
+# === Model & Data Settings ===
+model: "SD3"
+resolution: 768
+lambda_func: v
+optimized_reg_weight: ./LCFM/SD3_loss_v_MSE_DIV2k_neg_prompt.npy
+regularizer_weight: 0.5
+likelihood_weight_mode: reg_weight
+likelihood_steps: 15
+early_stopping: 1.e-4
+epochs: 1
+guidance: 2
+quantize: False
+use_tiny_ae: True
+negative_prompt: ""
+reg-shift: 0.0  
+
+# === Optimization Settings ===
+optimizer: 
+  name: SGD
+  kwargs:
+    lr: 1
+optimizer_dataterm: 
+  name: SGD
+  kwargs:
+    lr: 12
+
+# === Sampling & Misc ===
+t_sampling: descending
+n_steps: 50
+inv_alpha: 1-t
+quantize: False
+ts_min: 0.18
+projection: False
+seed: 3
+
+# === Experiment-Specific Settings ===
+prompt: a high quality photo of
+# Super-resolution x12 for FFHQ
+degradation:
+  name: SuperRes
+  kwargs:
+    scale: 12
+    noise_std: 0.0
+    img_size: 768
+likelihood_weight: 1

demo_images/demo_0_meta.json

@@ -0,0 +1,7 @@
+{
+    "prompt": "a high quality image of a face.",
+    "seed_on_slider": 2646030500,
+    "use_random_seed_checkbox": true,
+    "num_steps": 50,
+    "task_type": "Super Resolution"
+}

demo_images/demo_1_meta.json

@@ -0,0 +1,7 @@
+{
+    "prompt": "a house",
+    "seed_on_slider": 434714511,
+    "use_random_seed_checkbox": false,
+    "num_steps": 50,
+    "task_type": "Inpainting"
+}

demo_images/demo_2_meta.json

@@ -0,0 +1,7 @@
+{
+    "prompt": "a high quality image of a face.",
+    "seed_on_slider": 3211750901,
+    "use_random_seed_checkbox": false,
+    "num_steps": 50,
+    "task_type": "Inpainting"
+}

demo_images/demo_3_meta.json

@@ -0,0 +1,7 @@
+{
+    "prompt": "a girl with bright blue eyes ",
+    "seed_on_slider": 4257757956,
+    "use_random_seed_checkbox": true,
+    "num_steps": 50,
+    "task_type": "Inpainting"
+}

inference_scripts/run_image_inv.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2024 <Julius Erbach ETH Zurich>
+#
+# This file is part of the var_post_samp project and is licensed under the MIT License.
+# See the LICENSE file in the project root for more information.
+"""
+Usage:
+    python run_image_inv.py --config <config.yaml>
+"""
+
+import os
+import sys
+import time
+import csv
+import yaml
+import torch
+import random
+import click
+import numpy as np
+import tqdm
+import datetime
+import torchvision
+
+from flair.helper_functions import parse_click_context
+from flair.pipelines import model_loader
+from flair.utils import data_utils
+from flair import var_post_samp
+
+
+dtype = torch.bfloat16
+
+num_gpus = torch.cuda.device_count()
+if num_gpus > 0:
+    devices = [f"cuda:{i}" for i in range(num_gpus)]
+    primary_device = devices[0]
+    print(f"Using devices: {devices}")
+    print(f"Primary device for operations: {primary_device}")
+else:
+    print("No CUDA devices found. Using CPU.")
+    devices = ["cpu"]
+    primary_device = "cpu"
+
+
+@click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
+@click.option("--config", "config_file_arg", type=click.Path(exists=True), help="Path to the config file")
+@click.option("--target_file", type=click.Path(exists=True), help="Path to the target file or folder")
+@click.option("--result_folder", type=click.Path(file_okay=False, dir_okay=True, writable=True, resolve_path=True), help="Path to the output folder. It will be created if it doesn't exist.")
+@click.option("--mask_file", type=click.Path(exists=True), default=None, help="Path to the mask file npy. Optional used for image inpainting. True pixels are observed.")
+@click.pass_context
+def main(ctx, config_file_arg, target_file, result_folder, mask_file=None):
+    """Main entry point for image inversion and sampling.
+
+    The user must provide either a caption_file (with per-image captions) OR a single prompt for all images in the config YAML file.
+    """
+    with open(config_file_arg, "r") as f:
+        config = yaml.safe_load(f)
+    ctx = parse_click_context(ctx)
+    config.update(ctx)
+    # Read caption_file and prompt from config
+    caption_file = config.get("caption_file", None)
+    prompt = config.get("prompt", None)
+
+    # Enforce mutually exclusive caption_file or prompt
+    if (not caption_file and not prompt) or (caption_file and prompt):
+        raise ValueError("You must provide either 'caption_file' OR 'prompt' (not both) in the config file. See documentation.")
+
+    # wandb removed, so config_dict is just a copy
+    config_dict = dict(config)
+    torch.manual_seed(config["seed"])
+    np.random.seed(config["seed"])
+    random.seed(config["seed"])
+
+    # Use config values as-is (no to_absolute_path)
+    caption_file = caption_file if caption_file else None
+
+    guidance_img_iterator = data_utils.yield_images(
+        target_file, size=config["resolution"]
+    )
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    counter = 1
+    name = f'results_{config["model"]}_{config["degradation"]["name"]}_resolution_{config["resolution"]}_noise_{config["degradation"]["kwargs"]["noise_std"]}_{timestamp}'
+    candidate = os.path.join(name)
+    while os.path.exists(candidate):
+        candidate = os.path.join(f"{name}_{counter}")
+        counter += 1
+    output_folders = data_utils.generate_output_structure(
+        result_folder,
+        [
+            candidate,
+            f'input_{config["degradation"]["name"]}_resolution_{config["resolution"]}_noise_{config["degradation"]["kwargs"]["noise_std"]}',
+            f'target_{config["degradation"]["name"]}_resolution_{config["resolution"]}_noise_{config["degradation"]["kwargs"]["noise_std"]}',
+        ],
+    )
+    config_out = os.path.join(os.path.split(output_folders[0])[0], "config.yaml")
+    with open(config_out, "w") as f:
+        yaml.safe_dump(config_dict, f)
+
+    source_files = list(data_utils.find_files(target_file, ext="png"))
+    num_images = len(source_files)
+    print(f"Found {num_images} images.")
+
+    # Load captions
+    if caption_file:
+        captions = data_utils.load_captions_from_file(caption_file, user_prompt="")
+        if not captions:
+            sys.exit("Error: No captions were loaded from the provided caption file.")
+        if len(captions) != num_images:
+            print("Warning: Number of captions does not match number of images.")
+        prompts_in_order = [captions.get(os.path.basename(f), "") for f in source_files]
+    else:
+        # Use the single prompt for all images
+        prompts_in_order = [prompt for _ in range(num_images)]
+
+    if any(p == "" for p in prompts_in_order):
+        print("Warning: Some images might not have corresponding captions or prompt is empty.")
+    config["prompt"] = prompts_in_order
+
+    model, inp_kwargs = model_loader.load_model(config, device=devices)
+    if mask_file and config["degradation"]["name"] == "Inpainting":
+        config["degradation"]["kwargs"]["mask"] = mask_file
+    posterior_model = var_post_samp.VariationalPosterior(model, config)
+    guidance_img_iterator = data_utils.yield_images(
+        target_file, size=config["resolution"]
+    )
+    for idx, guidance_img in tqdm.tqdm(enumerate(guidance_img_iterator), total=num_images):
+        guidance_img = guidance_img.to(dtype).cuda()
+        y = posterior_model.forward_operator(guidance_img)
+        tic = time.time()
+        with torch.no_grad():
+            result_dict = posterior_model.forward(y, inp_kwargs[idx])
+            x_hat = result_dict["x_hat"]
+        toc = time.time()
+        print(f"Runtime: {toc - tic}")
+        guidance_img = guidance_img.cuda()
+        result_file = output_folders[0].format(idx)
+        input_file = output_folders[1].format(idx)
+        ground_truth_file = output_folders[2].format(idx)
+        x_hat_pil = torchvision.transforms.ToPILImage()(
+            x_hat.float()[0].clip(-1, 1) * 0.5 + 0.5
+        )
+        x_hat_pil.save(result_file)
+        try:
+            if config["degradation"]["name"] == "SuperRes":
+                input_img = posterior_model.forward_operator.nn(y)
+            else:
+                input_img = posterior_model.forward_operator.pseudo_inv(y)
+            input_img_pil = torchvision.transforms.ToPILImage()(
+                input_img.float()[0].clip(-1, 1) * 0.5 + 0.5
+            )
+            input_img_pil.save(input_file)
+        except Exception:
+            print("Error in pseudo-inverse operation. Skipping input image save.")
+        guidance_img_pil = torchvision.transforms.ToPILImage()(
+            guidance_img.float()[0] * 0.5 + 0.5
+        )
+        guidance_img_pil.save(ground_truth_file)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,18 @@
+tqdm
+torch
+click
+pyyaml
+numpy
+torchvision
+diffusers
+Pillow
+scipy
+munch
+transformers
+torchmetrics
+scikit-image
+opencv-python
+sentencepiece
+protobuf
+accelerate
+gradio

scripts/compute_metrics.py

@@ -0,0 +1,144 @@
+import torchmetrics
+import torch
+from PIL import Image
+import argparse
+from flair.utils import data_utils
+import os
+import tqdm
+import torch.nn.functional as F
+from torchmetrics.image.kid import KernelInceptionDistance
+
+
+MAX_BATCH_SIZE = None
+
+@torch.no_grad()
+def main(args):
+    # Determine device
+    if args.device == "cuda" and torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    print(f"Using device: {device}")
+
+    # load images
+    gt_iterator = data_utils.yield_images(os.path.abspath(args.gt), size=args.resolution)
+    pred_iterator = data_utils.yield_images(os.path.abspath(args.pred), size=args.resolution)
+    fid_metric = torchmetrics.image.fid.FrechetInceptionDistance(normalize=True).to(device)
+    # kid_metric = KernelInceptionDistance(subset_size=args.kid_subset_size, normalize=True).to(device)
+    lpips_metric = torchmetrics.image.lpip.LearnedPerceptualImagePatchSimilarity(
+        net_type="alex", normalize=False, reduction="mean"
+    ).to(device)
+    if args.patch_size:
+        patch_fid_metric = torchmetrics.image.fid.FrechetInceptionDistance(normalize=True).to(device)
+        # patch_kid_metric = KernelInceptionDistance(subset_size=args.kid_subset_size, normalize=True).to(device)
+    psnr_list = []
+    lpips_list = []
+    ssim_list = []
+    # iterate over images
+    for gt, pred in tqdm.tqdm(zip(gt_iterator, pred_iterator)):
+        # Move tensors to the selected device
+        gt = gt.to(device)
+        pred = pred.to(device)
+
+        # resize gt to pred size
+        if gt.shape[-2:] != (args.resolution, args.resolution):
+            gt = F.interpolate(gt, size=args.resolution, mode="area")
+        if pred.shape[-2:] != (args.resolution, args.resolution):
+            pred = F.interpolate(pred, size=args.resolution, mode="area")
+        # to range [0,1]
+        gt_norm = gt * 0.5 + 0.5
+        pred_norm = pred * 0.5 + 0.5
+        # compute PSNR
+        psnr = torchmetrics.functional.image.peak_signal_noise_ratio(
+            pred_norm, gt_norm, data_range=1.0
+        )
+        psnr_list.append(psnr.cpu()) # Move result to CPU
+        # compute LPIPS
+        lpips_score = lpips_metric(pred.clip(-1,1), gt.clip(-1,1))
+        lpips_list.append(lpips_score.cpu()) # Move result to CPU
+        # compute SSIM
+        ssim = torchmetrics.functional.image.structural_similarity_index_measure(
+            pred_norm, gt_norm, data_range=1.0
+        )
+        ssim_list.append(ssim.cpu()) # Move result to CPU
+        print(f"PSNR: {psnr}, LPIPS: {lpips_score}, SSIM: {ssim}")
+        # compute FID
+        # Ensure inputs are on the correct device (already handled by moving gt/pred earlier)
+        fid_metric.update(gt_norm, real=False)
+        fid_metric.update(pred_norm, real=True)
+        # compute KID
+        # kid_metric.update(pred, real=False)
+        # kid_metric.update(gt, real=True)
+        # compute Patchwise FID/KID if patch_size is specified
+        if args.patch_size:
+            # Extract patches
+            patch_size = args.patch_size
+            gt_patches = F.unfold(gt_norm, kernel_size=patch_size, stride=patch_size)
+            pred_patches = F.unfold(pred_norm, kernel_size=patch_size, stride=patch_size)
+            # Reshape patches: (B, C*P*P, N_patches) -> (B*N_patches, C, P, P)
+            B, C, H, W = gt.shape
+            N_patches = gt_patches.shape[-1]
+            gt_patches = gt_patches.permute(0, 2, 1).reshape(B * N_patches, C, patch_size, patch_size)
+            pred_patches = pred_patches.permute(0, 2, 1).reshape(B * N_patches, C, patch_size, patch_size)
+            # Update patch FID metric (inputs are already on the correct device)
+            # Update patch KID metric
+            # process mini batches of patches
+            if MAX_BATCH_SIZE is None:
+                patch_fid_metric.update(pred_patches, real=False)
+                patch_fid_metric.update(gt_patches, real=True)
+                # patch_kid_metric.update(pred_patches, real=False)
+                # patch_kid_metric.update(gt_patches, real=True)
+            else:
+                for i in range(0, N_patches, MAX_BATCH_SIZE):
+                    patch_fid_metric.update(pred_patches[i:i + MAX_BATCH_SIZE], real=False)
+                    patch_fid_metric.update(gt_patches[i:i + MAX_BATCH_SIZE], real=True)
+                    # patch_kid_metric.update(pred_patches[i:i + MAX_BATCH_SIZE], real=False)
+                    # patch_kid_metric.update(gt_patches[i:i + MAX_BATCH_SIZE], real=True)
+
+    # compute FID
+    fid = fid_metric.compute()
+    # compute KID
+    # kid_mean, kid_std = kid_metric.compute()
+    if args.patch_size:
+        patch_fid = patch_fid_metric.compute()
+        # patch_kid_mean, patch_kid_std = patch_kid_metric.compute()
+    # compute average metrics (on CPU)
+    avg_psnr = torch.mean(torch.stack(psnr_list))
+    avg_lpips = torch.mean(torch.stack(lpips_list))
+    avg_ssim = torch.mean(torch.stack(ssim_list))
+    # compute standard deviation (on CPU)
+    std_psnr = torch.std(torch.stack(psnr_list))
+    std_lpips = torch.std(torch.stack(lpips_list))
+    std_ssim = torch.std(torch.stack(ssim_list))
+    print(f"PSNR: {avg_psnr} +/- {std_psnr}")
+    print(f"LPIPS: {avg_lpips} +/- {std_lpips}")
+    print(f"SSIM: {avg_ssim} +/- {std_ssim}")
+    print(f"FID: {fid}") # FID is computed on the selected device, print directly
+    # print(f"KID: {kid_mean} +/- {kid_std}") # KID is computed on the selected device, print directly
+    if args.patch_size:
+        print(f"Patch FID ({args.patch_size}x{args.patch_size}): {patch_fid}") # Patch FID is computed on the selected device, print directly
+        # print(f"Patch KID ({args.patch_size}x{args.patch_size}): {patch_kid_mean} +/- {patch_kid_std}") # Patch KID is computed on the selected device, print directly
+    # save to prediction folder
+    out_file = os.path.join(args.pred, "fid_metrics.txt")
+    with open(out_file, "w") as f:
+        f.write(f"PSNR: {avg_psnr.item()} +/- {std_psnr.item()}\n") # Use .item() for scalar tensors
+        f.write(f"LPIPS: {avg_lpips.item()} +/- {std_lpips.item()}\n")
+        f.write(f"SSIM: {avg_ssim.item()} +/- {std_ssim.item()}\n")
+        f.write(f"FID: {fid.item()}\n") # Use .item() for scalar tensors
+        # f.write(f"KID: {kid_mean.item()} +/- {kid_std.item()}\n") # Use .item() for scalar tensors
+        if args.patch_size:
+            f.write(f"Patch FID ({args.patch_size}x{args.patch_size}): {patch_fid.item()}\n") # Use .item() for scalar tensors
+            # f.write(f"Patch KID ({args.patch_size}x{args.patch_size}): {patch_kid_mean.item()} +/- {patch_kid_std.item()}\n") # Use .item() for scalar tensors
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Compute metrics")
+    parser.add_argument("--gt", type=str, help="Path to ground truth image")
+    parser.add_argument("--pred", type=str, help="Path to predicted image")
+    parser.add_argument("--resolution", type=int, default=768, help="resolution at which to evaluate")
+    parser.add_argument("--patch_size", type=int, default=None, help="Patch size for Patchwise FID/KID computation (e.g., 12). If None, skip.")
+    parser.add_argument("--kid_subset_size", type=int, default=1000, help="Subset size for KID computation.")
+    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda"], help="Device to run computation on (cpu or cuda)")
+    args = parser.parse_args()
+
+    main(args)

scripts/generate_caption.py

@@ -0,0 +1,107 @@
+import sys
+import os
+import argparse
+from PIL import Image
+
+# Add the path to the thirdparty/SeeSR directory to the Python path
+sys.path.append(os.path.abspath("./thirdparty/SeeSR"))
+
+import torch
+from torchvision import transforms
+from ram.models.ram_lora import ram
+from ram import inference_ram as inference
+
+def load_ram_model(ram_model_path: str, dape_model_path: str):
+    """
+    Load the RAM model with the given paths.
+
+    Args:
+        ram_model_path (str): Path to the pretrained RAM model.
+        dape_model_path (str): Path to the pretrained DAPE model.
+
+    Returns:
+        torch.nn.Module: Loaded RAM model.
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load the RAM model
+    tag_model = ram(pretrained=ram_model_path, pretrained_condition=dape_model_path, image_size=384, vit="swin_l")
+    tag_model.eval()
+    return tag_model.to(device)
+
+def generate_caption(image_path: str, tag_model) -> str:
+    """
+    Generate a caption for a degraded image using the RAM model.
+
+    Args:
+        image_path (str): Path to the degraded input image.
+        tag_model (torch.nn.Module): Preloaded RAM model.
+
+    Returns:
+        str: Generated caption for the image.
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Define image transformations
+    tensor_transforms = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+    ram_transforms = transforms.Compose([
+        transforms.Resize((384, 384)),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    # Load and preprocess the image
+    image = Image.open(image_path).convert("RGB")
+    image_tensor = tensor_transforms(image).unsqueeze(0).to(device)
+    image_tensor = ram_transforms(image_tensor)
+
+    # Generate caption using the RAM model
+    caption = inference(image_tensor, tag_model)
+
+    return caption[0]
+
+def process_images_in_directory(input_dir: str, output_file: str, tag_model):
+    """
+    Process all images in a directory, generate captions using the RAM model,
+    and save the captions to a file.
+
+    Args:
+        input_dir (str): Path to the directory containing input images.
+        output_file (str): Path to the file where captions will be saved.
+        tag_model (torch.nn.Module): Preloaded RAM model.
+    """
+    # Open the output file for writing captions
+    with open(output_file, "w") as f:
+        # Iterate through all files in the input directory
+        for filename in os.listdir(input_dir):
+            # Construct the full path to the image file
+            image_path = os.path.join(input_dir, filename)
+
+            # Check if the file is an image
+            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
+                try:
+                    # Generate a caption for the image
+                    caption = generate_caption(image_path, tag_model)
+                    print(f"Generated caption for {filename}: {caption}")
+                    # Write the caption to the output file
+                    f.write(f"{filename}: {caption}\n")
+
+                    print(f"Processed {filename}: {caption}")
+                except Exception as e:
+                    print(f"Error processing {filename}: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate captions for images using RAM and DAPE models.")
+    parser.add_argument("--input_dir", type=str, default="data/val", help="Path to the directory containing input images.")
+    parser.add_argument("--output_file", type=str, default="data/val_captions.txt", help="Path to the file where captions will be saved.")
+    parser.add_argument("--ram_model", type=str, default="/home/erbachj/scratch2/projects/var_post_samp/thirdparty/SeeSR/preset/model/ram_swin_large_14m.pth", help="Path to the pretrained RAM model.")
+    parser.add_argument("--dape_model", type=str, default="/home/erbachj/scratch2/projects/var_post_samp/thirdparty/SeeSR/preset/model/DAPE.pth", help="Path to the pretrained DAPE model.")
+
+    args = parser.parse_args()
+
+    # Load the RAM model once
+    tag_model = load_ram_model(args.ram_model, args.dape_model)
+
+    # Process images in the directory
+    process_images_in_directory(args.input_dir, args.output_file, tag_model)

setup.py

@@ -0,0 +1,14 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="flair",
+    version="0.1.0",
+    author="Julius Erbach",
+    author_email="[email protected]",
+    description="Solving Inverse Problems with FLAIR",
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    include_package_data=True,
+    python_requires=">=3.7",
+    keywords="pytorch variational posterior sampling deep learning",
+)

src/flair/degradations.py

@@ -0,0 +1,198 @@
+import torch
+import numpy as np
+from munch import munchify
+from scipy.ndimage import distance_transform_edt
+from flair.functions.degradation import get_degradation
+import torchvision 
+
+class BaseDegradation(torch.nn.Module):
+    def __init__(self, noise_std=0.0):
+        super().__init__()
+        self.noise_std = noise_std
+
+    def forward(self, x):
+        x = x + self.noise_std * torch.randn_like(x)
+        return x
+
+    def pseudo_inv(self, y):
+        return y
+
+
+def zero_filler(x, scale):
+    B, C, H, W = x.shape
+    scale = int(scale)
+    H_new, W_new = H * scale, W * scale
+    out = torch.zeros(B, C, H_new, W_new, dtype=x.dtype, device=x.device)
+    out[:, :, ::scale, ::scale] = x
+    return out
+
+
+class SuperRes(BaseDegradation):
+    def __init__(self, scale, noise_std=0.0, img_size=256):
+        super().__init__(noise_std=noise_std)
+        self.scale = scale
+        deg_config = munchify({
+        'channels': 3,
+        'image_size': img_size,
+        'deg_scale': scale
+        })
+        self.img_size = img_size
+        self.deg = get_degradation("sr_bicubic", deg_config, device="cuda")
+
+    def forward(self, x, noise=True):
+        dtype = x.dtype
+        y = self.deg.A(x.float())
+        # add noise
+        if noise:
+            y = super().forward(y)
+
+        return y.to(dtype)
+
+    def pseudo_inv(self, y):
+        x = self.deg.At(y.float()).reshape(1,3,self.img_size, self.img_size)* self.scale**2
+        return x.to(y.dtype)
+    
+    
+    def nn(self, y):
+        x = torch.nn.functional.interpolate(
+            y.reshape(1,3,self.img_size//self.scale, self.img_size//self.scale), scale_factor=self.scale, mode="nearest"
+        )
+        return x.to(y.dtype)
+    
+class SuperResGradio(BaseDegradation):
+    def __init__(self, scale, noise_std=0.0, img_size=256):
+        super().__init__(noise_std=noise_std)
+        self.scale = scale
+        self.downscaler = lambda x: torch.nn.functional.interpolate(
+            x.float(), scale_factor=1/self.scale, mode="bilinear", align_corners=False, antialias=True
+        )
+        self.upscaler = lambda x: torch.nn.functional.interpolate(
+            x.float(), scale_factor=self.scale, mode="bilinear", align_corners=False, antialias=True
+        )
+        self.img_size = img_size 
+
+    def forward(self, x, noise=True):
+        dtype = x.dtype
+        y = self.downscaler(x.float())
+        # add noise
+        if noise:
+            y = super().forward(y)
+
+        return y.to(dtype)
+
+    def pseudo_inv(self, y):
+        x = self.upscaler(y.float())
+        return x.to(y.dtype)
+    
+    
+    def nn(self, y):
+        x = torch.nn.functional.interpolate(
+            y.reshape(1,3,self.img_size//self.scale, self.img_size//self.scale), scale_factor=self.scale, mode="nearest-exact"
+        )
+        return x.to(y.dtype)
+
+
+class Inpainting(BaseDegradation):
+    def __init__(self, mask, H, W, noise_std=0.0):
+        """
+        mask: torch.Tensor, shape (H, W), dtype bool
+        function assumes 3 channels
+        """
+        super().__init__(noise_std=noise_std)
+        if isinstance(mask, list):
+            # generate box from left, right, lower upper list
+            # observed region is True
+            mask_ = torch.ones(H, W, dtype=torch.bool)
+            mask_[slice(*mask[0:2]), slice(*mask[2:])] = False
+            # repeat for 3 channels
+            mask_ = mask_.repeat(3, 1, 1)
+        elif isinstance(mask, str):
+            # load mask file
+            mask_ = torch.tensor(np.load(mask), dtype=torch.bool)
+            mask_ = mask_.repeat(3, 1, 1)
+        elif isinstance(mask, torch.Tensor):
+            if mask.ndim == 2:
+                # assume mask is for one channel, repeat for 3 channels
+                mask_ = mask[None].repeat(3, 1, 1)
+            elif mask.ndim == 3 and mask.shape[0] == 1:
+                # assume mask is for one channel, repeat for 3 channels
+                mask_ = mask.repeat(3, 1, 1)
+            else:
+                mask_ = mask
+        else:
+            raise ValueError("Mask must be a list, string (file path), or torch.Tensor.")
+        self.mask = mask_
+        self.H, self.W = H, W
+
+    def forward(self, x, noise=True):
+        B = x.shape[0]
+        y = x[self.mask[None]].view(B, -1)
+        # add noise
+        if noise:
+            y = super().forward(y)
+        return y
+
+    def pseudo_inv(self, y):
+        x = torch.zeros(y.shape[0], 3 * self.H * self.W, dtype=y.dtype, device=y.device)
+        x[:, self.mask.view(-1)] = y
+        x = x.view(y.shape[0], 3, self.H, self.W)
+        # x = inpaint_nearest(x[0], self.mask[0])[None]
+        return x
+
+
+def inpaint_nearest(image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    Fill missing pixels in an image using the nearest observed pixel value.
+
+    Args:
+        image: A tensor of shape [C, H, W] representing the image.
+        mask: A tensor of shape [H, W] with 1 for observed pixels and 0 for missing pixels.
+
+    Returns:
+        A tensor of shape [C, H, W] where missing pixels have been filled.
+    """
+    # Move tensors to CPU and convert to numpy arrays.
+    image_np = image.cpu().float().numpy()
+    # Convert mask to boolean: True for observed, False for missing.
+    mask_np = mask.cpu().numpy().astype(bool)
+
+    # Compute the distance transform of the inverse mask (~mask_np).
+    # The function returns:
+    #   - distances: distance to the nearest True pixel in mask_np
+    #   - indices: the indices of that nearest True pixel for each pixel.
+    # indices has shape (2, H, W): first row is the row index, second row is the column index.
+    _, indices = distance_transform_edt(~mask_np, return_indices=True)
+
+    # Create a copy of the image to hold the filled values.
+    filled_image_np = np.empty_like(image_np)
+
+    # For each channel, replace every pixel with the value of the nearest observed pixel.
+    for c in range(image_np.shape[0]):
+        filled_image_np[c] = image_np[c, indices[0], indices[1]]
+
+    # Convert back to a torch tensor and send to the original device.
+    return torch.from_numpy(filled_image_np).to(image.device).to(image.dtype)
+    
+class MotionBlur(BaseDegradation):
+    def __init__(self, kernel_size=5, img_size=256, noise_std=0.0):
+        super().__init__(noise_std=noise_std)
+        deg_config = munchify({
+            'channels': 3,
+            'image_size': img_size,
+            'deg_scale': kernel_size
+        })
+        self.img_size = img_size
+        self.deg = get_degradation("deblur_motion", deg_config, device="cuda")
+
+    def forward(self, x, noise=True):
+        dtype = x.dtype
+        y = self.deg.A(x.float())
+        # add noise
+        if noise:
+            y = super().forward(y)
+        return y.to(dtype)
+    
+    def pseudo_inv(self, y):
+        dtype = y.dtype
+        x = self.deg.At(y.float()).reshape(1,3,self.img_size, self.img_size)
+        return x.to(dtype)

src/flair/functions/ckpt_util.py

@@ -0,0 +1,72 @@
+import os, hashlib
+import requests
+from tqdm import tqdm
+
+URL_MAP = {
+    "cifar10": "https://heibox.uni-heidelberg.de/f/869980b53bf5416c8a28/?dl=1",
+    "ema_cifar10": "https://heibox.uni-heidelberg.de/f/2e4f01e2d9ee49bab1d5/?dl=1",
+    "lsun_bedroom": "https://heibox.uni-heidelberg.de/f/f179d4f21ebc4d43bbfe/?dl=1",
+    "ema_lsun_bedroom": "https://heibox.uni-heidelberg.de/f/b95206528f384185889b/?dl=1",
+    "lsun_cat": "https://heibox.uni-heidelberg.de/f/fac870bd988348eab88e/?dl=1",
+    "ema_lsun_cat": "https://heibox.uni-heidelberg.de/f/0701aac3aa69457bbe34/?dl=1",
+    "lsun_church": "https://heibox.uni-heidelberg.de/f/2711a6f712e34b06b9d8/?dl=1",
+    "ema_lsun_church": "https://heibox.uni-heidelberg.de/f/44ccb50ef3c6436db52e/?dl=1",
+}
+CKPT_MAP = {
+    "cifar10": "diffusion_cifar10_model/model-790000.ckpt",
+    "ema_cifar10": "ema_diffusion_cifar10_model/model-790000.ckpt",
+    "lsun_bedroom": "diffusion_lsun_bedroom_model/model-2388000.ckpt",
+    "ema_lsun_bedroom": "ema_diffusion_lsun_bedroom_model/model-2388000.ckpt",
+    "lsun_cat": "diffusion_lsun_cat_model/model-1761000.ckpt",
+    "ema_lsun_cat": "ema_diffusion_lsun_cat_model/model-1761000.ckpt",
+    "lsun_church": "diffusion_lsun_church_model/model-4432000.ckpt",
+    "ema_lsun_church": "ema_diffusion_lsun_church_model/model-4432000.ckpt",
+}
+MD5_MAP = {
+    "cifar10": "82ed3067fd1002f5cf4c339fb80c4669",
+    "ema_cifar10": "1fa350b952534ae442b1d5235cce5cd3",
+    "lsun_bedroom": "f70280ac0e08b8e696f42cb8e948ff1c",
+    "ema_lsun_bedroom": "1921fa46b66a3665e450e42f36c2720f",
+    "lsun_cat": "bbee0e7c3d7abfb6e2539eaf2fb9987b",
+    "ema_lsun_cat": "646f23f4821f2459b8bafc57fd824558",
+    "lsun_church": "eb619b8a5ab95ef80f94ce8a5488dae3",
+    "ema_lsun_church": "fdc68a23938c2397caba4a260bc2445f",
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root=None, check=False, prefix='exp'):
+    if 'church_outdoor' in name:
+        name = name.replace('church_outdoor', 'church')
+    assert name in URL_MAP
+    # Modify the path when necessary
+    cachedir = os.environ.get("XDG_CACHE_HOME", os.path.join(prefix, "logs/"))
+    root = (
+        root
+        if root is not None
+        else os.path.join(cachedir, "diffusion_models_converted")
+    )
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path

src/flair/functions/conjugate_gradient.py

@@ -0,0 +1,66 @@
+# From wikipedia. MATLAB code,
+# function x = conjgrad(A, b, x)
+#     r = b - A * x;
+#     p = r;
+#     rsold = r' * r;
+# 
+#     for i = 1:length(b)
+#         Ap = A * p;
+#         alpha = rsold / (p' * Ap);
+#         x = x + alpha * p;
+#         r = r - alpha * Ap;
+#         rsnew = r' * r;
+#         if sqrt(rsnew) < 1e-10
+#             break
+#         end
+#         p = r + (rsnew / rsold) * p;
+#         rsold = rsnew;
+#     end
+# end
+
+from typing import Callable, Optional
+
+import torch
+
+
+def CG(A: Callable,
+       b: torch.Tensor,
+       x: torch.Tensor,
+       m: Optional[int]=5,
+       eps: Optional[float]=1e-4,
+       damping: float=0.0,
+       use_mm: bool=False) -> torch.Tensor:
+    
+    if use_mm:
+        mm_fn = lambda x, y: torch.mm(x.view(1, -1), y.view(1, -1).T)
+    else:
+        mm_fn = lambda x, y: (x * y).flatten().sum()
+    
+    orig_shape = x.shape
+    x = x.view(x.shape[0], -1)
+
+    r = b - A(x)
+    p = r.clone()
+
+    rsold = mm_fn(r, r)
+    assert not (rsold != rsold).any(), print(f'NaN detected 1')
+
+    for i in range(m):
+        Ap = A(p) + damping * p
+        alpha = rsold / mm_fn(p, Ap)
+        assert not (alpha != alpha).any(), print(f'NaN detected 2')
+        
+        x = x + alpha * p
+        r = r - alpha * Ap
+
+        rsnew = mm_fn(r, r)
+        assert not (rsnew != rsnew).any(), print('NaN detected 3')
+        
+        if rsnew.sqrt().abs() < eps:
+            break
+
+        p = r + (rsnew / rsold) * p
+        rsold = rsnew.clone()
+
+    return x.reshape(orig_shape)
+

src/flair/functions/degradation.py

@@ -0,0 +1,211 @@
+import numpy as np
+import torch
+from munch import Munch
+
+from flair.functions import svd_operators as svd_op
+from flair.functions import measurements
+from flair.utils.inpaint_util import MaskGenerator
+
+__DEGRADATION__ = {}
+
+def register_degradation(name: str):
+    def wrapper(fn):
+        if __DEGRADATION__.get(name) is not None:
+            raise NameError(f'DEGRADATION {name} is already registered')
+        __DEGRADATION__[name]=fn
+        return fn
+    return wrapper
+
+def get_degradation(name: str,
+                    deg_config: Munch,
+                    device:torch.device):
+    if __DEGRADATION__.get(name) is None:
+        raise NameError(f'DEGRADATION {name} does not exist.')
+    return __DEGRADATION__[name](deg_config, device)
+
+@register_degradation(name='cs_walshhadamard')
+def deg_cs_walshhadamard(deg_config, device):
+    compressed_size = round(1/deg_config.deg_scale)
+    A_funcs = svd_op.WalshHadamardCS(deg_config.channels,
+                                     deg_config.image_size,
+                                     compressed_size,
+                                     torch.randperm(deg_config.image_size**2),
+                                     device)
+    return A_funcs
+
+@register_degradation(name='cs_blockbased')
+def deg_cs_blockbased(deg_config, device):
+    cs_ratio = deg_config.deg_scale
+    A_funcs = svd_op.CS(deg_config.channels,
+                        deg_config.image_size,
+                        cs_ratio,
+                        device)
+    return A_funcs
+
+@register_degradation(name='inpainting')
+def deg_inpainting(deg_config, device):
+    # TODO: generate mask rather than load
+    loaded = np.load("exp/inp_masks/mask_768_half.npy")  # block
+    # loaded = np.load("lip_mask_4.npy")
+    mask = torch.from_numpy(loaded).to(device).reshape(-1)
+    missing_r = torch.nonzero(mask == 0).long().reshape(-1) * 3
+    missing_g = missing_r + 1
+    missing_b = missing_g + 1
+    missing = torch.cat([missing_r, missing_g, missing_b], dim=0)
+    A_funcs = svd_op.Inpainting(deg_config.channels,
+                                deg_config.image_size,
+                                missing,
+                                device)
+    return A_funcs
+
+@register_degradation(name='denoising')
+def deg_denoise(deg_config, device):
+    A_funcs = svd_op.Denoising(deg_config.channels,
+                               deg_config.image_size,
+                               device)
+    return A_funcs
+
+@register_degradation(name='colorization')
+def deg_colorization(deg_config, device):
+    A_funcs = svd_op.Colorization(deg_config.image_size,
+                                  device)
+    return A_funcs
+
+
+@register_degradation(name='sr_avgpool')
+def deg_sr_avgpool(deg_config, device):
+    blur_by = int(deg_config.deg_scale)
+    A_funcs = svd_op.SuperResolution(deg_config.channels,
+                                     deg_config.image_size,
+                                     blur_by,
+                                     device)
+    return A_funcs
+
+@register_degradation(name='sr_bicubic')
+def deg_sr_bicubic(deg_config, device):
+    def bicubic_kernel(x, a=-0.5):
+        if abs(x) <= 1:
+            return (a + 2) * abs(x) ** 3 - (a + 3) * abs(x) ** 2 + 1
+        elif 1 < abs(x) and abs(x) < 2:
+            return a * abs(x) ** 3 - 5 * a * abs(x) ** 2 + 8 * a * abs(x) - 4 * a
+        else:
+            return 0
+
+    factor = int(deg_config.deg_scale)
+    k = np.zeros((factor * 4))
+    for i in range(factor * 4):
+        x = (1 / factor) * (i - np.floor(factor * 4 / 2) + 0.5)
+        k[i] = bicubic_kernel(x)
+    k = k / np.sum(k)
+    kernel = torch.from_numpy(k).float().to(device)
+    A_funcs = svd_op.SRConv(kernel / kernel.sum(),
+                            deg_config.channels,
+                            deg_config.image_size,
+                            device,
+                            stride=factor)
+    return A_funcs
+
+@register_degradation(name='deblur_uni')
+def deg_deblur_uni(deg_config, device):
+    A_funcs = svd_op.Deblurring(torch.tensor([1/deg_config.deg_scale]*deg_config.deg_scale).to(device),
+                                deg_config.channels,
+                                deg_config.image_size,
+                                device)
+    return A_funcs
+
+@register_degradation(name='deblur_gauss')
+def deg_deblur_gauss(deg_config, device):
+    sigma = 3.0
+    pdf = lambda x: torch.exp(torch.Tensor([-0.5 * (x / sigma) ** 2]))
+    size = deg_config.deg_scale
+    ker = []
+    for k in range(-size//2, size//2):
+        ker.append(pdf(k))
+    kernel = torch.Tensor(ker).to(device)
+    A_funcs = svd_op.Deblurring(kernel / kernel.sum(),
+                                deg_config.channels,
+                                deg_config.image_size,
+                                device)
+    return A_funcs
+
+@register_degradation(name='deblur_aniso')
+def deg_deblur_aniso(deg_config, device):
+    sigma = 20
+    pdf = lambda x: torch.exp(torch.Tensor([-0.5 * (x / sigma) ** 2]))
+    kernel2 = torch.Tensor([pdf(-4), pdf(-3), pdf(-2), pdf(-1), pdf(0), pdf(1), pdf(2), pdf(3), pdf(4)]).to(device)
+
+    sigma = 1
+    pdf = lambda x: torch.exp(torch.Tensor([-0.5 * (x / sigma) ** 2]))
+    kernel1 = torch.Tensor([pdf(-4), pdf(-3), pdf(-2), pdf(-1), pdf(0), pdf(1), pdf(2), pdf(3), pdf(4)]).to(device)
+
+    A_funcs = svd_op.Deblurring2D(kernel1 / kernel1.sum(),
+                                  kernel2 / kernel2.sum(),
+                                  deg_config.channels,
+                                  deg_config.image_size,
+                                  device)
+    return A_funcs
+
+@register_degradation(name='deblur_motion')
+def deg_deblur_motion(deg_config, device):
+    A_funcs = measurements.MotionBlurOperator(
+        kernel_size=deg_config.deg_scale,
+        intensity=0.5,
+        device=device
+    )
+    return A_funcs
+
+@register_degradation(name='deblur_nonuniform')
+def deg_deblur_motion(deg_config, device, kernels=None, masks=None):
+    A_funcs = measurements.NonuniformBlurOperator(
+        deg_config.image_size,
+        deg_config.deg_scale,
+        device,
+        kernels=kernels,
+        masks=masks,
+    )
+    return A_funcs
+
+
+# ======= FOR arbitraty image size =======
+@register_degradation(name='sr_avgpool_gen')
+def deg_sr_avgpool_general(deg_config, device):
+    blur_by = int(deg_config.deg_scale)
+    A_funcs = svd_op.SuperResolutionGeneral(deg_config.channels,
+                                            deg_config.imgH,
+                                            deg_config.imgW,
+                                            blur_by,
+                                            device)
+    return A_funcs
+
+@register_degradation(name='deblur_gauss_gen')
+def deg_deblur_guass_general(deg_config, device):
+    A_funcs = measurements.GaussialBlurOperator(
+        kernel_size=deg_config.deg_scale,
+        intensity=3.0,
+        device=device
+    )
+    return A_funcs
+
+
+from flair.functions.jpeg import jpeg_encode, jpeg_decode
+
+class JPEGOperator():
+    def __init__(self, qf: int, device):
+        self.qf = qf
+        self.device = device
+
+    def A(self, img):
+        x_luma, x_chroma = jpeg_encode(img, self.qf)
+        return x_luma, x_chroma
+
+    def At(self, encoded):
+        return jpeg_decode(encoded, self.qf)
+
+
+@register_degradation(name='jpeg')
+def deg_jpeg(deg_config, device):
+    A_funcs = JPEGOperator(
+        qf = deg_config.deg_scale,
+        device=device
+    )
+    return A_funcs

src/flair/functions/jpeg.py

@@ -0,0 +1,392 @@
+# ---------------------------------------------------------------
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+#
+# This file has been modified from ddrm-jpeg.
+#
+# Source:
+# https://github.com/bahjat-kawar/ddrm-jpeg/blob/master/functions/jpeg_torch.py
+#
+# The license for the original version of this file can be
+# found in this directory (LICENSE_DDRM_JPEG).
+# The modifications to this file are subject to the same license.
+# ---------------------------------------------------------------
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+def dct1(x):
+    """
+    Discrete Cosine Transform, Type I
+    :param x: the input signal
+    :return: the DCT-I of the signal over the last dimension
+    """
+    x_shape = x.shape
+    x = x.view(-1, x_shape[-1])
+
+    return torch.fft.rfft(torch.cat([x, x.flip([1])[:, 1:-1]], dim=1))[:, :, 0].view(*x_shape)
+
+
+def idct1(X):
+    """
+    The inverse of DCT-I, which is just a scaled DCT-I
+    Our definition if idct1 is such that idct1(dct1(x)) == x
+    :param X: the input signal
+    :return: the inverse DCT-I of the signal over the last dimension
+    """
+    n = X.shape[-1]
+    return dct1(X) / (2 * (n - 1))
+
+
+def dct(x, norm=None):
+    """
+    Discrete Cosine Transform, Type II (a.k.a. the DCT)
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last dimension
+    """
+    x_shape = x.shape
+    N = x_shape[-1]
+    x = x.contiguous().view(-1, N)
+
+    v = torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim=1)
+
+    Vc = torch.view_as_real(torch.fft.fft(v, dim=1))
+
+    k = - torch.arange(N, dtype=x.dtype, device=x.device)[None, :] * np.pi / (2 * N)
+    W_r = torch.cos(k)
+    W_i = torch.sin(k)
+
+    V = Vc[:, :, 0] * W_r - Vc[:, :, 1] * W_i
+
+    if norm == 'ortho':
+        V[:, 0] /= np.sqrt(N) * 2
+        V[:, 1:] /= np.sqrt(N / 2) * 2
+
+    V = 2 * V.view(*x_shape)
+
+    return V
+
+
+def idct(X, norm=None):
+    """
+    The inverse to DCT-II, which is a scaled Discrete Cosine Transform, Type III
+    Our definition of idct is that idct(dct(x)) == x
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the inverse DCT-II of the signal over the last dimension
+    """
+
+    x_shape = X.shape
+    N = x_shape[-1]
+
+    X_v = X.contiguous().view(-1, x_shape[-1]) / 2
+
+    if norm == 'ortho':
+        X_v[:, 0] *= np.sqrt(N) * 2
+        X_v[:, 1:] *= np.sqrt(N / 2) * 2
+
+    k = torch.arange(x_shape[-1], dtype=X.dtype, device=X.device)[None, :] * np.pi / (2 * N)
+    W_r = torch.cos(k)
+    W_i = torch.sin(k)
+
+    V_t_r = X_v
+    V_t_i = torch.cat([X_v[:, :1] * 0, -X_v.flip([1])[:, :-1]], dim=1)
+
+    V_r = V_t_r * W_r - V_t_i * W_i
+    V_i = V_t_r * W_i + V_t_i * W_r
+
+    V = torch.cat([V_r.unsqueeze(2), V_i.unsqueeze(2)], dim=2)
+
+    v = torch.fft.irfft(torch.view_as_complex(V), n=V.shape[1], dim=1)
+    x = v.new_zeros(v.shape)
+    x[:, ::2] += v[:, :N - (N // 2)]
+    x[:, 1::2] += v.flip([1])[:, :N // 2]
+
+    return x.view(*x_shape)
+
+
+def dct_2d(x, norm=None):
+    """
+    2-dimentional Discrete Cosine Transform, Type II (a.k.a. the DCT)
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 2 dimensions
+    """
+    X1 = dct(x, norm=norm)
+    X2 = dct(X1.transpose(-1, -2), norm=norm)
+    return X2.transpose(-1, -2)
+
+
+def idct_2d(X, norm=None):
+    """
+    The inverse to 2D DCT-II, which is a scaled Discrete Cosine Transform, Type III
+    Our definition of idct is that idct_2d(dct_2d(x)) == x
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 2 dimensions
+    """
+    x1 = idct(X, norm=norm)
+    x2 = idct(x1.transpose(-1, -2), norm=norm)
+    return x2.transpose(-1, -2)
+
+
+def dct_3d(x, norm=None):
+    """
+    3-dimentional Discrete Cosine Transform, Type II (a.k.a. the DCT)
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 3 dimensions
+    """
+    X1 = dct(x, norm=norm)
+    X2 = dct(X1.transpose(-1, -2), norm=norm)
+    X3 = dct(X2.transpose(-1, -3), norm=norm)
+    return X3.transpose(-1, -3).transpose(-1, -2)
+
+
+def idct_3d(X, norm=None):
+    """
+    The inverse to 3D DCT-II, which is a scaled Discrete Cosine Transform, Type III
+    Our definition of idct is that idct_3d(dct_3d(x)) == x
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 3 dimensions
+    """
+    x1 = idct(X, norm=norm)
+    x2 = idct(x1.transpose(-1, -2), norm=norm)
+    x3 = idct(x2.transpose(-1, -3), norm=norm)
+    return x3.transpose(-1, -3).transpose(-1, -2)
+
+
+class LinearDCT(nn.Linear):
+    """Implement any DCT as a linear layer; in practice this executes around
+    50x faster on GPU. Unfortunately, the DCT matrix is stored, which will
+    increase memory usage.
+    :param in_features: size of expected input
+    :param type: which dct function in this file to use"""
+    def __init__(self, in_features, type, norm=None, bias=False):
+        self.type = type
+        self.N = in_features
+        self.norm = norm
+        super(LinearDCT, self).__init__(in_features, in_features, bias=bias)
+
+    def reset_parameters(self):
+        # initialise using dct function
+        I = torch.eye(self.N)
+        if self.type == 'dct1':
+            self.weight.data = dct1(I).data.t()
+        elif self.type == 'idct1':
+            self.weight.data = idct1(I).data.t()
+        elif self.type == 'dct':
+            self.weight.data = dct(I, norm=self.norm).data.t()
+        elif self.type == 'idct':
+            self.weight.data = idct(I, norm=self.norm).data.t()
+        self.weight.requires_grad = False # don't learn this!
+
+
+def apply_linear_2d(x, linear_layer):
+    """Can be used with a LinearDCT layer to do a 2D DCT.
+    :param x: the input signal
+    :param linear_layer: any PyTorch Linear layer
+    :return: result of linear layer applied to last 2 dimensions
+    """
+    X1 = linear_layer(x)
+    X2 = linear_layer(X1.transpose(-1, -2))
+    return X2.transpose(-1, -2)
+
+
+def apply_linear_3d(x, linear_layer):
+    """Can be used with a LinearDCT layer to do a 3D DCT.
+    :param x: the input signal
+    :param linear_layer: any PyTorch Linear layer
+    :return: result of linear layer applied to last 3 dimensions
+    """
+    X1 = linear_layer(x)
+    X2 = linear_layer(X1.transpose(-1, -2))
+    X3 = linear_layer(X2.transpose(-1, -3))
+    return X3.transpose(-1, -3).transpose(-1, -2)
+
+
+def torch_rgb2ycbcr(x):
+    # Assume x is a batch of size (N x C x H x W)
+    v = torch.tensor([[.299, .587, .114], [-.1687, -.3313, .5], [.5, -.4187, -.0813]]).to(x.device)
+    ycbcr = torch.tensordot(x, v, dims=([1], [1])).transpose(3, 2).transpose(2, 1)
+    ycbcr[:,1:] += 128
+    return ycbcr
+
+
+def torch_ycbcr2rgb(x):
+    # Assume x is a batch of size (N x C x H x W)
+    v = torch.tensor([[ 1.00000000e+00, -3.68199903e-05,  1.40198758e+00],
+       [ 1.00000000e+00, -3.44113281e-01, -7.14103821e-01],
+       [ 1.00000000e+00,  1.77197812e+00, -1.34583413e-04]]).to(x.device)
+    x[:, 1:] -= 128
+    rgb = torch.tensordot(x, v, dims=([1], [1])).transpose(3, 2).transpose(2, 1)
+    return rgb
+
+def chroma_subsample(x):
+    return x[:, 0:1, :, :], x[:, 1:, ::2, ::2]
+
+
+def general_quant_matrix(qf = 10):
+    q1 = torch.tensor([
+    16,  11,  10,  16,  24,  40,  51,  61,
+    12,  12,  14,  19,  26,  58,  60,  55,
+    14,  13,  16,  24,  40,  57,  69,  56,
+    14,  17,  22,  29,  51,  87,  80,  62,
+    18,  22,  37,  56,  68, 109, 103,  77,
+    24,  35,  55,  64,  81, 104, 113,  92,
+    49,  64,  78,  87, 103, 121, 120, 101,
+    72,  92,  95,  98, 112, 100, 103,  99
+    ])
+    q2 = torch.tensor([
+        17,  18,  24,  47,  99,  99,  99,  99,
+        18,  21,  26,  66,  99,  99,  99,  99,
+        24,  26,  56,  99,  99,  99,  99,  99,
+        47,  66,  99,  99,  99,  99,  99,  99,
+        99,  99,  99,  99,  99,  99,  99,  99,
+        99,  99,  99,  99,  99,  99,  99,  99,
+        99,  99,  99,  99,  99,  99,  99,  99,
+        99,  99,  99,  99,  99,  99,  99,  99
+    ])
+    s = (5000 / qf) if qf < 50 else (200 - 2 * qf)
+    q1 = torch.floor((s * q1 + 50) / 100)
+    q1[q1 <= 0] = 1
+    q1[q1 > 255] = 255
+    q2 = torch.floor((s * q2 + 50) / 100)
+    q2[q2 <= 0] = 1
+    q2[q2 > 255] = 255
+    return q1, q2
+
+
+def quantization_matrix(qf):
+    return general_quant_matrix(qf)
+    # q1 = torch.tensor([[ 80,  55,  50,  80, 120, 200, 255, 255],
+    #                    [ 60,  60,  70,  95, 130, 255, 255, 255],
+    #                    [ 70,  65,  80, 120, 200, 255, 255, 255],
+    #                    [ 70,  85, 110, 145, 255, 255, 255, 255],
+    #                    [ 90, 110, 185, 255, 255, 255, 255, 255],
+    #                    [120, 175, 255, 255, 255, 255, 255, 255],
+    #                    [245, 255, 255, 255, 255, 255, 255, 255],
+    #                    [255, 255, 255, 255, 255, 255, 255, 255]])
+    # q2 = torch.tensor([[ 85,  90, 120, 235, 255, 255, 255, 255],
+    #                    [ 90, 105, 130, 255, 255, 255, 255, 255],
+    #                    [120, 130, 255, 255, 255, 255, 255, 255],
+    #                    [235, 255, 255, 255, 255, 255, 255, 255],
+    #                    [255, 255, 255, 255, 255, 255, 255, 255],
+    #                    [255, 255, 255, 255, 255, 255, 255, 255],
+    #                    [255, 255, 255, 255, 255, 255, 255, 255],
+    #                    [255, 255, 255, 255, 255, 255, 255, 255]])
+    # return q1, q2
+
+def jpeg_encode(x, qf):
+    # Assume x is a batch of size (N x C x H x W)
+    # [-1, 1] to [0, 255]
+    x = (x + 1) / 2 * 255
+    n_batch, _, n_size, _ = x.shape
+
+    x = torch_rgb2ycbcr(x)
+    x_luma, x_chroma = chroma_subsample(x)
+    unfold = nn.Unfold(kernel_size=(8, 8), stride=(8, 8))
+    x_luma = unfold(x_luma).transpose(2, 1)
+    x_chroma = unfold(x_chroma).transpose(2, 1)
+
+    x_luma = x_luma.reshape(-1, 8, 8) - 128
+    x_chroma = x_chroma.reshape(-1, 8, 8) - 128
+
+    dct_layer = LinearDCT(8, 'dct', norm='ortho')
+    dct_layer.to(x_luma.device)
+    x_luma = apply_linear_2d(x_luma, dct_layer)
+    x_chroma = apply_linear_2d(x_chroma, dct_layer)
+
+    x_luma = x_luma.view(-1, 1, 8, 8)
+    x_chroma = x_chroma.view(-1, 2, 8, 8)
+
+    q1, q2 = quantization_matrix(qf)
+    q1 = q1.to(x_luma.device)
+    q2 = q2.to(x_luma.device)
+    x_luma /= q1.view(1, 8, 8)
+    x_chroma /= q2.view(1, 8, 8)
+
+    x_luma = x_luma.round()
+    x_chroma = x_chroma.round()
+
+    x_luma = x_luma.reshape(n_batch, (n_size // 8) ** 2, 64).transpose(2, 1)
+    x_chroma = x_chroma.reshape(n_batch, (n_size // 16) ** 2, 64 * 2).transpose(2, 1)
+
+    fold = nn.Fold(output_size=(n_size, n_size), kernel_size=(8, 8), stride=(8, 8))
+    x_luma = fold(x_luma)
+    fold = nn.Fold(output_size=(n_size // 2, n_size // 2), kernel_size=(8, 8), stride=(8, 8))
+    x_chroma = fold(x_chroma)
+
+    return [x_luma, x_chroma]
+
+
+
+def jpeg_decode(x, qf):
+    # Assume x[0] is a batch of size (N x 1 x H x W) (luma)
+    # Assume x[1:] is a batch of size (N x 2 x H/2 x W/2) (chroma)
+    x_luma, x_chroma = x
+    n_batch, _, n_size, _ = x_luma.shape
+    unfold = nn.Unfold(kernel_size=(8, 8), stride=(8, 8))
+    x_luma = unfold(x_luma).transpose(2, 1)
+    x_luma = x_luma.reshape(-1, 1, 8, 8)
+    x_chroma = unfold(x_chroma).transpose(2, 1)
+    x_chroma = x_chroma.reshape(-1, 2, 8, 8)
+
+    q1, q2 = quantization_matrix(qf)
+    q1 = q1.to(x_luma.device)
+    q2 = q2.to(x_luma.device)
+    x_luma *= q1.view(1, 8, 8)
+    x_chroma *= q2.view(1, 8, 8)
+
+    x_luma = x_luma.reshape(-1, 8, 8)
+    x_chroma = x_chroma.reshape(-1, 8, 8)
+
+    dct_layer = LinearDCT(8, 'idct', norm='ortho')
+    dct_layer.to(x_luma.device)
+    x_luma = apply_linear_2d(x_luma, dct_layer)
+    x_chroma = apply_linear_2d(x_chroma, dct_layer)
+
+    x_luma = (x_luma + 128).reshape(n_batch, (n_size // 8) ** 2, 64).transpose(2, 1)
+    x_chroma = (x_chroma + 128).reshape(n_batch, (n_size // 16) ** 2, 64 * 2).transpose(2, 1)
+
+    fold = nn.Fold(output_size=(n_size, n_size), kernel_size=(8, 8), stride=(8, 8))
+    x_luma = fold(x_luma)
+    fold = nn.Fold(output_size=(n_size // 2, n_size // 2), kernel_size=(8, 8), stride=(8, 8))
+    x_chroma = fold(x_chroma)
+
+    x_chroma_repeated = torch.zeros(n_batch, 2, n_size, n_size, device = x_luma.device)
+    x_chroma_repeated[:, :, 0::2, 0::2] = x_chroma
+    x_chroma_repeated[:, :, 0::2, 1::2] = x_chroma
+    x_chroma_repeated[:, :, 1::2, 0::2] = x_chroma
+    x_chroma_repeated[:, :, 1::2, 1::2] = x_chroma
+
+    x = torch.cat([x_luma, x_chroma_repeated], dim=1)
+
+    x = torch_ycbcr2rgb(x)
+
+    # [0, 255] to [-1, 1]
+    x = x / 255 * 2 - 1
+
+    return x
+
+
+def build_jpeg(qf):
+    # log.info(f"[Corrupt] JPEG restoration: {qf=}  ...")
+    def jpeg(img):
+        encoded = jpeg_encode(img, qf)
+        return jpeg_decode(encoded, qf), encoded
+    return jpeg

src/flair/functions/measurements.py

@@ -0,0 +1,429 @@
+'''This module handles task-dependent operations (A) and noises (n) to simulate a measurement y=Ax+n.'''
+
+from abc import ABC, abstractmethod
+from functools import partial
+
+from torch.nn import functional as F
+from torchvision import torch
+
+from flair.utils.blur_util import Blurkernel
+from flair.utils.img_util import fft2d
+import numpy as np
+from flair.utils.resizer import Resizer
+from flair.utils.utils_sisr import pre_calculate_FK, pre_calculate_nonuniform
+from torch.fft import fft2, ifft2
+
+
+from flair.motionblur.motionblur import Kernel
+
+# =================
+# Operation classes
+# =================
+
+__OPERATOR__ = {}
+_GAMMA_FACTOR = 2.2
+
+def register_operator(name: str):
+    def wrapper(cls):
+        if __OPERATOR__.get(name, None):
+            raise NameError(f"Name {name} is already registered!")
+        __OPERATOR__[name] = cls
+        return cls
+    return wrapper
+
+
+def get_operator(name: str, **kwargs):
+    if __OPERATOR__.get(name, None) is None:
+        raise NameError(f"Name {name} is not defined.")
+    return __OPERATOR__[name](**kwargs)
+
+
+class LinearOperator(ABC):
+    @abstractmethod
+    def forward(self, data, **kwargs):
+        # calculate A * X
+        pass
+
+    @abstractmethod
+    def noisy_forward(self, data, **kwargs):
+        # calculate A * X + n
+        pass
+
+    @abstractmethod
+    def transpose(self, data, **kwargs):
+        # calculate A^T * X
+        pass
+
+    def ortho_project(self, data, **kwargs):
+        # calculate (I - A^T * A)X
+        return data - self.transpose(self.forward(data, **kwargs), **kwargs)
+
+    def project(self, data, measurement, **kwargs):
+        # calculate (I - A^T * A)Y - AX
+        return self.ortho_project(measurement, **kwargs) - self.forward(data, **kwargs)
+
+
+@register_operator(name='noise')
+class DenoiseOperator(LinearOperator):
+    def __init__(self, device):
+        self.device = device
+
+    def forward(self, data):
+        return data
+
+    def noisy_forward(self, data):
+        return data
+
+    def transpose(self, data):
+        return data
+
+    def ortho_project(self, data):
+        return data
+
+    def project(self, data):
+        return data
+
+
+@register_operator(name='sr_bicubic')
+class SuperResolutionOperator(LinearOperator):
+    def __init__(self,
+                 in_shape,
+                 scale_factor,
+                 noise,
+                 noise_scale,
+                 device):
+        self.device = device
+        self.up_sample = partial(F.interpolate, scale_factor=scale_factor)
+        self.down_sample = Resizer(in_shape, 1/scale_factor).to(device)
+        self.noise = get_noise(name=noise, scale=noise_scale)
+
+    def A(self, data, **kwargs):
+        return self.forward(data, **kwargs)
+
+    def forward(self, data, **kwargs):
+        return self.down_sample(data)
+
+    def noisy_forward(self, data, **kwargs):
+        return self.noise.forward(self.down_sample(data))
+
+    def transpose(self, data, **kwargs):
+        return self.up_sample(data)
+
+    def project(self, data, measurement, **kwargs):
+        return data - self.transpose(self.forward(data)) + self.transpose(measurement)
+
+@register_operator(name='deblur_motion')
+class MotionBlurOperator(LinearOperator):
+    def __init__(self,
+                 kernel_size,
+                 intensity,
+                 device):
+        self.device = device
+        self.kernel_size = kernel_size
+        self.conv = Blurkernel(blur_type='motion',
+                               kernel_size=kernel_size,
+                               std=intensity,
+                               device=device).to(device)  # should we keep this device term?
+
+        self.kernel_size =kernel_size
+        self.intensity = intensity
+        self.kernel = Kernel(size=(kernel_size, kernel_size), intensity=intensity)
+        kernel = torch.tensor(self.kernel.kernelMatrix, dtype=torch.float32)
+        self.conv.update_weights(kernel)
+
+    def forward(self, data, **kwargs):
+        # A^T * A
+        return self.conv(data)
+
+    def noisy_forward(self, data, **kwargs):
+        pass
+
+    def transpose(self, data, **kwargs):
+        return data
+
+    def change_kernel(self):
+        self.kernel = Kernel(size=(self.kernel_size, self.kernel_size), intensity=self.intensity)
+        kernel = torch.tensor(self.kernel.kernelMatrix, dtype=torch.float32)
+        self.conv.update_weights(kernel)
+
+    def get_kernel(self):
+        kernel = self.kernel.kernelMatrix.type(torch.float32).to(self.device)
+        return kernel.view(1, 1, self.kernel_size, self.kernel_size)
+
+    def A(self, data):
+        return self.forward(data)
+
+    def At(self, data):
+        return self.transpose(data)
+
+@register_operator(name='deblur_gauss')
+class GaussialBlurOperator(LinearOperator):
+    def __init__(self,
+                 kernel_size,
+                 intensity,
+                 device):
+        self.device = device
+        self.kernel_size = kernel_size
+        self.conv = Blurkernel(blur_type='gaussian',
+                               kernel_size=kernel_size,
+                               std=intensity,
+                               device=device).to(device)
+        self.kernel = self.conv.get_kernel()
+        self.conv.update_weights(self.kernel.type(torch.float32))
+
+    def forward(self, data, **kwargs):
+        return self.conv(data)
+
+    def noisy_forward(self, data, **kwargs):
+        pass
+
+    def transpose(self, data, **kwargs):
+        return data
+
+    def get_kernel(self):
+        return self.kernel.view(1, 1, self.kernel_size, self.kernel_size)
+
+    def apply_kernel(self, data, kernel):
+        self.conv.update_weights(kernel.type(torch.float32))
+        return self.conv(data)
+
+    def A(self, data):
+        return self.forward(data)
+
+    def At(self, data):
+        return self.transpose(data)
+
+@register_operator(name='inpainting')
+class InpaintingOperator(LinearOperator):
+    '''This operator get pre-defined mask and return masked image.'''
+    def __init__(self,
+                 noise,
+                 noise_scale,
+                 device):
+        self.device = device
+        self.noise = get_noise(name=noise, scale=noise_scale)
+
+    def forward(self, data, **kwargs):
+        try:
+            return data * kwargs.get('mask', None).to(self.device)
+        except:
+            raise ValueError("Require mask")
+
+    def noisy_forward(self, data, **kwargs):
+        return self.noise.forward(self.forward(data, **kwargs))
+
+    def transpose(self, data, **kwargs):
+        return data
+
+    def ortho_project(self, data, **kwargs):
+        return data - self.forward(data, **kwargs)
+
+# Operator for BlindDPS.
+@register_operator(name='blind_blur')
+class BlindBlurOperator(LinearOperator):
+    def __init__(self, device, **kwargs) -> None:
+        self.device = device
+
+    def forward(self, data, kernel, **kwargs):
+        return self.apply_kernel(data, kernel)
+
+    def transpose(self, data, **kwargs):
+        return data
+
+    def apply_kernel(self, data, kernel):
+        #TODO: faster way to apply conv?:W
+
+        b_img = torch.zeros_like(data).to(self.device)
+        for i in range(3):
+            b_img[:, i, :, :] = F.conv2d(data[:, i:i+1, :, :], kernel, padding='same')
+        return b_img
+
+
+class NonLinearOperator(ABC):
+    @abstractmethod
+    def forward(self, data, **kwargs):
+        pass
+
+    @abstractmethod
+    def noisy_forward(self, data, **kwargs):
+        pass
+
+    def project(self, data, measurement, **kwargs):
+        return data + measurement - self.forward(data)
+
+@register_operator(name='phase_retrieval')
+class PhaseRetrievalOperator(NonLinearOperator):
+    def __init__(self,
+                 oversample,
+                 noise,
+                 noise_scale,
+                 device):
+        self.pad = int((oversample / 8.0) * 256)
+        self.device = device
+        self.noise = get_noise(name=noise, scale=noise_scale)
+
+    def forward(self, data, **kwargs):
+        padded = F.pad(data, (self.pad, self.pad, self.pad, self.pad))
+        amplitude = fft2d(padded).abs()
+        return amplitude
+
+    def noisy_forard(self, data, **kwargs):
+        return self.noise.forward(self.forward(data, **kwargs))
+
+@register_operator(name='nonuniform_blur')
+class NonuniformBlurOperator(LinearOperator):
+    def __init__(self, in_shape, kernel_size, device,
+                 kernels=None, masks=None):
+        self.device = device
+        self.kernel_size = kernel_size
+        self.in_shape = in_shape
+
+        # TODO: generalize
+        if kernels is None and masks is None:
+            self.kernels = np.load('./functions/nonuniform/kernels/000001.npy')
+            self.masks = np.load('./functions/nonuniform/masks/000001.npy')
+            self.kernels = torch.tensor(self.kernels).to(device)
+            self.masks = torch.tensor(self.masks).to(device)
+
+    # approximate in image space
+    def forward_img(self, data):
+        K = self.kernel_size
+        data = F.pad(data, (K//2, K//2, K//2, K//2), mode="reflect")
+        kernels = self.kernels.transpose(0, 1)
+        data_rgb_batch = data.transpose(0, 1)
+        conv_rgb_batch = F.conv2d(data_rgb_batch, kernels)
+        y_rgb_batch = conv_rgb_batch * self.masks
+        y_rgb_batch = torch.sum(y_rgb_batch, dim=1, keepdim=True)
+        y = y_rgb_batch.transpose(0, 1)
+        return y
+
+    # NOTE: Only using this operator will make the problem nonlinear (gamma-correction)
+    def forward_nonlinear(self, data, flatten=False, noiseless=False):
+        # 1. Usual nonuniform blurring degradataion pipeline
+        kernels = self.kernels.transpose(0, 1)
+        FK, FKC = pre_calculate_FK(kernels)
+        y = ifft2(FK * fft2(data)).real
+        y = y.transpose(0, 1)
+        y_rgb_batch = self.masks * y
+        y_rgb_batch = torch.sum(y_rgb_batch, dim=1, keepdim=True)
+        y = y_rgb_batch.transpose(0, 1)
+        F2KM, FKFMy = pre_calculate_nonuniform(data, y, FK, FKC, self.masks)
+        self.pre_calculated = (FK, FKC, F2KM, FKFMy)
+        # 2. Gamma-correction
+        y = (y + 1) / 2
+        y = y ** (1 / _GAMMA_FACTOR)
+        y = (y - 0.5) / 0.5
+        return y
+
+    def noisy_forward(self, data, **kwargs):
+        return self.noise.forward(self.forward(data))
+
+    # exact in Fourier
+    def forward(self, data, flatten=False, noiseless=False):
+        # [1, 25, 33, 33] -> [25, 1, 33, 33]
+        kernels = self.kernels.transpose(0, 1)
+        # [25, 1, 512, 512]
+        FK, FKC = pre_calculate_FK(kernels)
+        # [25, 3, 512, 512]
+        y = ifft2(FK * fft2(data)).real
+        # [3, 25, 512, 512]
+        y = y.transpose(0, 1)
+        y_rgb_batch = self.masks * y
+        # [3, 1, 512, 512]
+        y_rgb_batch = torch.sum(y_rgb_batch, dim=1, keepdim=True)
+        # [1, 3, 512, 512]
+        y = y_rgb_batch.transpose(0, 1)
+        F2KM, FKFMy = pre_calculate_nonuniform(data, y, FK, FKC, self.masks)
+        self.pre_calculated = (FK, FKC, F2KM, FKFMy)
+        return y
+
+    def transpose(self, y, flatten=False):
+        kernels = self.kernels.transpose(0, 1)
+        FK, FKC = pre_calculate_FK(kernels)
+        # 1. braodcast and multiply
+        # [3, 1, 512, 512]
+        y_rgb_batch = y.transpose(0, 1)
+        # [3, 25, 512, 512]
+        y_rgb_batch = y_rgb_batch.repeat(1, 25, 1, 1)
+        y = self.masks * y_rgb_batch
+        # 2. transpose of convolution in Fourier
+        # [25, 3, 512, 512]
+        y = y.transpose(0, 1)
+        ATy_broadcast = ifft2(FKC * fft2(y)).real
+        # [1, 3, 512, 512]
+        ATy = torch.sum(ATy_broadcast, dim=0, keepdim=True)
+        return ATy
+
+    def A(self, data):
+        return self.forward(data)
+
+    def At(self, data):
+        return self.transpose(data)
+
+# =============
+# Noise classes
+# =============
+
+
+__NOISE__ = {}
+
+def register_noise(name: str):
+    def wrapper(cls):
+        if __NOISE__.get(name, None):
+            raise NameError(f"Name {name} is already defined!")
+        __NOISE__[name] = cls
+        return cls
+    return wrapper
+
+def get_noise(name: str, **kwargs):
+    if __NOISE__.get(name, None) is None:
+        raise NameError(f"Name {name} is not defined.")
+    noiser = __NOISE__[name](**kwargs)
+    noiser.__name__ = name
+    return noiser
+
+class Noise(ABC):
+    def __call__(self, data):
+        return self.forward(data)
+
+    @abstractmethod
+    def forward(self, data):
+        pass
+
+@register_noise(name='clean')
+class Clean(Noise):
+    def __init__(self, **kwargs):
+        pass
+
+    def forward(self, data):
+        return data
+
+@register_noise(name='gaussian')
+class GaussianNoise(Noise):
+    def __init__(self, scale):
+        self.scale = scale
+
+    def forward(self, data):
+        return data + torch.randn_like(data, device=data.device) * self.scale
+
+
+@register_noise(name='poisson')
+class PoissonNoise(Noise):
+    def __init__(self, scale):
+        self.scale = scale
+
+    def forward(self, data):
+        '''
+        Follow skimage.util.random_noise.
+        '''
+
+        # version 3 (stack-overflow)
+        import numpy as np
+        data = (data + 1.0) / 2.0
+        data = data.clamp(0, 1)
+        device = data.device
+        data = data.detach().cpu()
+        data = torch.from_numpy(np.random.poisson(data * 255.0 * self.scale) / 255.0 / self.scale)
+        data = data * 2.0 - 1.0
+        data = data.clamp(-1, 1)
+        return data.to(device)

src/flair/functions/nonuniform/kernels/000001.npy

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

src/flair/functions/svd_ddnm.py

@@ -0,0 +1,206 @@
+import torch
+from tqdm import tqdm
+import torchvision.utils as tvu
+import torchvision
+import os
+
+class_num = 951
+
+
+def compute_alpha(beta, t):
+    beta = torch.cat([torch.zeros(1).to(beta.device), beta], dim=0)
+    a = (1 - beta).cumprod(dim=0).index_select(0, t + 1).view(-1, 1, 1, 1)
+    return a
+
+def inverse_data_transform(x):
+    x = (x + 1.0) / 2.0
+    return torch.clamp(x, 0.0, 1.0)
+
+def ddnm_diffusion(x, model, b, eta, A_funcs, y, cls_fn=None, classes=None, config=None):
+    with torch.no_grad():
+
+        # setup iteration variables
+        skip = config.diffusion.num_diffusion_timesteps//config.time_travel.T_sampling
+        n = x.size(0)
+        x0_preds = []
+        xs = [x]
+
+        # generate time schedule
+        times = get_schedule_jump(config.time_travel.T_sampling, 
+                               config.time_travel.travel_length, 
+                               config.time_travel.travel_repeat,
+                              )
+        time_pairs = list(zip(times[:-1], times[1:]))
+        
+        # reverse diffusion sampling
+        for i, j in tqdm(time_pairs):
+            i, j = i*skip, j*skip
+            if j<0: j=-1 
+
+            if j < i: # normal sampling 
+                t = (torch.ones(n) * i).to(x.device)
+                next_t = (torch.ones(n) * j).to(x.device)
+                at = compute_alpha(b, t.long())
+                at_next = compute_alpha(b, next_t.long())
+                xt = xs[-1].to('cuda')
+                if cls_fn == None:
+                    et = model(xt, t)
+                else:
+                    classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
+                    et = model(xt, t, classes)
+                    et = et[:, :3]
+                    et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)
+
+                if et.size(1) == 6:
+                    et = et[:, :3]
+
+                x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
+
+                x0_t_hat = x0_t - A_funcs.A_pinv(
+                    A_funcs.A(x0_t.reshape(x0_t.size(0), -1)) - y.reshape(y.size(0), -1)
+                ).reshape(*x0_t.size())
+
+                c1 = (1 - at_next).sqrt() * eta
+                c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
+                xt_next = at_next.sqrt() * x0_t_hat + c1 * torch.randn_like(x0_t) + c2 * et
+
+                x0_preds.append(x0_t.to('cpu'))
+                xs.append(xt_next.to('cpu'))
+            else: # time-travel back
+                next_t = (torch.ones(n) * j).to(x.device)
+                at_next = compute_alpha(b, next_t.long())
+                x0_t = x0_preds[-1].to('cuda')
+                
+                xt_next = at_next.sqrt() * x0_t + torch.randn_like(x0_t) * (1 - at_next).sqrt()
+
+                xs.append(xt_next.to('cpu'))
+
+    return [xs[-1]], [x0_preds[-1]]
+
+def ddnm_plus_diffusion(x, model, b, eta, A_funcs, y, sigma_y, cls_fn=None, classes=None, config=None):
+    with torch.no_grad():
+
+        # setup iteration variables
+        skip = config.diffusion.num_diffusion_timesteps//config.time_travel.T_sampling
+        n = x.size(0)
+        x0_preds = []
+        xs = [x]
+
+        # generate time schedule
+        times = get_schedule_jump(config.time_travel.T_sampling, 
+                               config.time_travel.travel_length, 
+                               config.time_travel.travel_repeat,
+                              )
+        time_pairs = list(zip(times[:-1], times[1:]))        
+        
+        # reverse diffusion sampling
+        for i, j in tqdm(time_pairs):
+            i, j = i*skip, j*skip
+            if j<0: j=-1 
+
+            if j < i: # normal sampling 
+                t = (torch.ones(n) * i).to(x.device)
+                next_t = (torch.ones(n) * j).to(x.device)
+                at = compute_alpha(b, t.long())
+                at_next = compute_alpha(b, next_t.long())
+                xt = xs[-1].to('cuda')
+                if cls_fn == None:
+                    et = model(xt, t)
+                else:
+                    classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
+                    et = model(xt, t, classes)
+                    et = et[:, :3]
+                    et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)
+
+                if et.size(1) == 6:
+                    et = et[:, :3]
+
+                # Eq. 12
+                x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
+
+                sigma_t = (1 - at_next).sqrt()[0, 0, 0, 0]
+
+                # Eq. 17
+                x0_t_hat = x0_t - A_funcs.Lambda(A_funcs.A_pinv(
+                    A_funcs.A(x0_t.reshape(x0_t.size(0), -1)) - y.reshape(y.size(0), -1)
+                ).reshape(x0_t.size(0), -1), at_next.sqrt()[0, 0, 0, 0], sigma_y, sigma_t, eta).reshape(*x0_t.size())
+
+                # Eq. 51
+                xt_next = at_next.sqrt() * x0_t_hat + A_funcs.Lambda_noise(
+                    torch.randn_like(x0_t).reshape(x0_t.size(0), -1), 
+                    at_next.sqrt()[0, 0, 0, 0], sigma_y, sigma_t, eta, et.reshape(et.size(0), -1)).reshape(*x0_t.size())
+
+                x0_preds.append(x0_t.to('cpu'))
+                xs.append(xt_next.to('cpu'))
+            else: # time-travel back
+                next_t = (torch.ones(n) * j).to(x.device)
+                at_next = compute_alpha(b, next_t.long())
+                x0_t = x0_preds[-1].to('cuda')
+                
+                xt_next = at_next.sqrt() * x0_t + torch.randn_like(x0_t) * (1 - at_next).sqrt()
+
+                xs.append(xt_next.to('cpu'))
+                
+#             #ablation
+#             if i%50==0:
+#                 os.makedirs('/userhome/wyh/ddnm/debug/x0t', exist_ok=True)
+#                 tvu.save_image(
+#                     inverse_data_transform(x0_t[0]),
+#                     os.path.join('/userhome/wyh/ddnm/debug/x0t', f"x0_t_{i}.png")
+#                 )
+                
+#                 os.makedirs('/userhome/wyh/ddnm/debug/x0_t_hat', exist_ok=True)
+#                 tvu.save_image(
+#                     inverse_data_transform(x0_t_hat[0]),
+#                     os.path.join('/userhome/wyh/ddnm/debug/x0_t_hat', f"x0_t_hat_{i}.png")
+#                 )
+                
+#                 os.makedirs('/userhome/wyh/ddnm/debug/xt_next', exist_ok=True)
+#                 tvu.save_image(
+#                     inverse_data_transform(xt_next[0]),
+#                     os.path.join('/userhome/wyh/ddnm/debug/xt_next', f"xt_next_{i}.png")
+#                 )
+
+    return [xs[-1]], [x0_preds[-1]]
+
+# form RePaint
+def get_schedule_jump(T_sampling, travel_length, travel_repeat):
+
+    jumps = {}
+    for j in range(0, T_sampling - travel_length, travel_length):
+        jumps[j] = travel_repeat - 1
+
+    t = T_sampling
+    ts = []
+
+    while t >= 1:
+        t = t-1
+        ts.append(t)
+
+        if jumps.get(t, 0) > 0:
+            jumps[t] = jumps[t] - 1
+            for _ in range(travel_length):
+                t = t + 1
+                ts.append(t)
+
+    ts.append(-1)
+
+    _check_times(ts, -1, T_sampling)
+
+    return ts
+
+def _check_times(times, t_0, T_sampling):
+    # Check end
+    assert times[0] > times[1], (times[0], times[1])
+
+    # Check beginning
+    assert times[-1] == -1, times[-1]
+
+    # Steplength = 1
+    for t_last, t_cur in zip(times[:-1], times[1:]):
+        assert abs(t_last - t_cur) == 1, (t_last, t_cur)
+
+    # Value range
+    for t in times:
+        assert t >= t_0, (t, t_0)
+        assert t <= T_sampling, (t, T_sampling)

src/flair/functions/svd_operators.py

@@ -0,0 +1,1308 @@
+import torch
+
+
+class A_functions:
+    """
+    A class replacing the SVD of a matrix A, perhaps efficiently.
+    All input vectors are of shape (Batch, ...).
+    All output vectors are of shape (Batch, DataDimension).
+    """
+
+    def V(self, vec):
+        """
+        Multiplies the input vector by V
+        """
+        raise NotImplementedError()
+
+    def Vt(self, vec):
+        """
+        Multiplies the input vector by V transposed
+        """
+        raise NotImplementedError()
+
+    def U(self, vec):
+        """
+        Multiplies the input vector by U
+        """
+        raise NotImplementedError()
+
+    def Ut(self, vec):
+        """
+        Multiplies the input vector by U transposed
+        """
+        raise NotImplementedError()
+
+    def singulars(self):
+        """
+        Returns a vector containing the singular values. The shape of the vector should be the same as the smaller dimension (like U)
+        """
+        raise NotImplementedError()
+
+    def add_zeros(self, vec):
+        """
+        Adds trailing zeros to turn a vector from the small dimension (U) to the big dimension (V)
+        """
+        raise NotImplementedError()
+
+    def A(self, vec):
+        """
+        Multiplies the input vector by A
+        """
+        temp = self.Vt(vec)
+        singulars = self.singulars()
+        return self.U(singulars * temp[:, :singulars.shape[0]])
+
+    def At(self, vec):
+        """
+        Multiplies the input vector by A transposed
+        """
+        temp = self.Ut(vec)
+        singulars = self.singulars()
+        return self.V(self.add_zeros(singulars * temp[:, :singulars.shape[0]]))
+
+    def A_pinv(self, vec):
+        """
+        Multiplies the input vector by the pseudo inverse of A
+        """
+        temp = self.Ut(vec)
+        singulars = self.singulars()
+
+        factors = 1. / singulars
+        factors[singulars == 0] = 0.
+
+#         temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] / singulars
+        temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] * factors
+        return self.V(self.add_zeros(temp))
+
+    def A_pinv_eta(self, vec, eta):
+        """
+        Multiplies the input vector by the pseudo inverse of A with factor eta
+        """
+        temp = self.Ut(vec)
+        singulars = self.singulars()
+        factors = singulars / (singulars*singulars+eta)
+#         print(temp.size(), factors.size(), singulars.size())
+        temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] * factors
+        return self.V(self.add_zeros(temp))
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        raise NotImplementedError()
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        raise NotImplementedError()
+
+
+# block-wise CS
+class CS(A_functions):
+    def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
+        self.img_dim = img_dim
+        self.channels = channels
+        self.y_dim = img_dim // 32
+        self.ratio = 32
+        A = torch.randn(32**2, 32**2).to(device)
+        _, _, self.V_small = torch.svd(A, some=False)
+        self.Vt_small = self.V_small.transpose(0, 1)
+        self.singulars_small = torch.ones(int(32 * 32 * ratio), device=device)
+        self.cs_size = self.singulars_small.size(0)
+
+    def V(self, vec):
+        #reorder the vector back into patches (because singulars are ordered descendingly)
+
+        temp = vec.clone().reshape(vec.shape[0], -1)
+        patches = torch.zeros(vec.size(0), self.channels * self.y_dim ** 2, self.ratio ** 2, device=vec.device)
+        patches[:, :, :self.cs_size] = temp[:, :self.channels * self.y_dim ** 2 * self.cs_size].contiguous().reshape(
+            vec.size(0), -1, self.cs_size)
+        patches[:, :, self.cs_size:] = temp[:, self.channels * self.y_dim ** 2 * self.cs_size:].contiguous().reshape(
+            vec.size(0), self.channels * self.y_dim ** 2, -1)
+
+        #multiply each patch by the small V
+        patches = torch.matmul(self.V_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #repatch the patches into an image
+        patches_orig = patches.reshape(vec.shape[0], self.channels, self.y_dim, self.y_dim, self.ratio, self.ratio)
+        recon = patches_orig.permute(0, 1, 2, 4, 3, 5).contiguous()
+        recon = recon.reshape(vec.shape[0], self.channels * self.img_dim ** 2)
+        return recon
+
+    def Vt(self, vec):
+        #extract flattened patches
+        patches = vec.clone().reshape(vec.shape[0], self.channels, self.img_dim, self.img_dim)
+        patches = patches.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches = patches.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #multiply each by the small V transposed
+        patches = torch.matmul(self.Vt_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #reorder the vector to have the first entry first (because singulars are ordered descendingly)
+        recon = torch.zeros(vec.shape[0], self.channels * self.img_dim**2, device=vec.device)
+        recon[:, :self.channels * self.y_dim ** 2 * self.cs_size] = patches[:, :, :, :self.cs_size].contiguous().reshape(
+            vec.shape[0], -1)
+        recon[:, self.channels * self.y_dim ** 2 * self.cs_size:] = patches[:, :, :, self.cs_size:].contiguous().reshape(
+            vec.shape[0], -1)
+        return recon
+
+    def U(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec): #U is 1x1, so U^T = U
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self.singulars_small.repeat(self.channels * self.y_dim**2)
+
+    def add_zeros(self, vec):
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp = torch.zeros((vec.shape[0], self.channels * self.img_dim**2), device=vec.device)
+        temp[:, :reshaped.shape[1]] = reshaped
+        return temp
+
+
+def color2gray(x):
+    x = x[:, 0:1, :, :] * 0.3333 + x[:, 1:2, :, :] * 0.3334 + x[:, 2:, :, :] * 0.3333
+    return x
+
+
+def gray2color(x):
+    base = 0.3333 ** 2 + 0.3334 ** 2 + 0.3333 ** 2
+    return torch.stack((x * 0.3333 / base, x * 0.3334 / base, x * 0.3333 / base), 1)
+
+
+#a memory inefficient implementation for any general degradation A
+class GeneralA(A_functions):
+    def mat_by_vec(self, M, v):
+        vshape = v.shape[1]
+        if len(v.shape) > 2: vshape = vshape * v.shape[2]
+        if len(v.shape) > 3: vshape = vshape * v.shape[3]
+        return torch.matmul(M, v.view(v.shape[0], vshape,
+                        1)).view(v.shape[0], M.shape[0])
+
+    def __init__(self, A):
+        self._U, self._singulars, self._V = torch.svd(A, some=False)
+        self._Vt = self._V.transpose(0, 1)
+        self._Ut = self._U.transpose(0, 1)
+
+        ZERO = 1e-3
+        self._singulars[self._singulars < ZERO] = 0
+        print(len([x.item() for x in self._singulars if x == 0]))
+
+    def V(self, vec):
+        return self.mat_by_vec(self._V, vec.clone())
+
+    def Vt(self, vec):
+        return self.mat_by_vec(self._Vt, vec.clone())
+
+    def U(self, vec):
+        return self.mat_by_vec(self._U, vec.clone())
+
+    def Ut(self, vec):
+        return self.mat_by_vec(self._Ut, vec.clone())
+
+    def singulars(self):
+        return self._singulars
+
+    def add_zeros(self, vec):
+        out = torch.zeros(vec.shape[0], self._V.shape[0], device=vec.device)
+        out[:, :self._U.shape[0]] = vec.clone().reshape(vec.shape[0], -1)
+        return out
+
+#Walsh-Hadamard Compressive Sensing
+class WalshHadamardCS(A_functions):
+    def fwht(self, vec): #the Fast Walsh Hadamard Transform is the same as its inverse
+        a = vec.reshape(vec.shape[0], self.channels, self.img_dim**2)
+        h = 1
+        while h < self.img_dim**2:
+            a = a.reshape(vec.shape[0], self.channels, -1, h * 2)
+            b = a.clone()
+            a[:, :, :, :h] = b[:, :, :, :h] + b[:, :, :, h:2*h]
+            a[:, :, :, h:2*h] = b[:, :, :, :h] - b[:, :, :, h:2*h]
+            h *= 2
+        a = a.reshape(vec.shape[0], self.channels, self.img_dim**2) / self.img_dim
+        return a
+
+    def __init__(self, channels, img_dim, ratio, perm, device):
+        self.channels = channels
+        self.img_dim = img_dim
+        self.ratio = ratio
+        self.perm = perm
+        self._singulars = torch.ones(channels * img_dim**2 // ratio, device=device)
+
+    def V(self, vec):
+        temp = torch.zeros(vec.shape[0], self.channels, self.img_dim**2, device=vec.device)
+        temp[:, :, self.perm] = vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        return self.fwht(temp).reshape(vec.shape[0], -1)
+
+    def Vt(self, vec):
+        return self.fwht(vec.clone())[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars
+
+    def add_zeros(self, vec):
+        out = torch.zeros(vec.shape[0], self.channels * self.img_dim**2, device=vec.device)
+        out[:, :self.channels * self.img_dim**2 // self.ratio] = vec.clone().reshape(vec.shape[0], -1)
+        return out
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        temp_vec = self.fwht(vec.clone())[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        singulars = self._singulars
+        lambda_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device)
+        temp = torch.zeros(self.channels * self.img_dim ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (
+                    singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, -1)
+        temp_vec = temp_vec * lambda_t
+
+        temp_out = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out[:, :, self.perm] = temp_vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        return self.fwht(temp_out).reshape(vec.shape[0], -1)
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        temp_vec = vec.clone().reshape(
+            vec.shape[0], self.channels, self.img_dim ** 2)[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+        temp_eps = epsilon.clone().reshape(
+            vec.shape[0], self.channels, self.img_dim ** 2)[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        d1_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        singulars = self._singulars
+        temp = torch.zeros(self.channels * self.img_dim ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(
+                change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, -1)
+        d2_t = d2_t.reshape(1, -1)
+
+        temp_vec = temp_vec * d1_t
+        temp_eps = temp_eps * d2_t
+
+        temp_out_vec = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out_vec[:, :, self.perm] = temp_vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        temp_out_vec = self.fwht(temp_out_vec).reshape(vec.shape[0], -1)
+
+        temp_out_eps = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out_eps[:, :, self.perm] = temp_eps.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        temp_out_eps = self.fwht(temp_out_eps).reshape(vec.shape[0], -1)
+
+        return temp_out_vec + temp_out_eps
+
+
+#Inpainting
+class Inpainting(A_functions):
+    def __init__(self, channels, img_dim, missing_indices, device):
+        self.channels = channels
+        self.img_dim = img_dim
+        self._singulars = torch.ones(channels * img_dim**2 - missing_indices.shape[0]).to(device)
+        self.missing_indices = missing_indices
+        self.kept_indices = torch.Tensor([i for i in range(channels * img_dim**2) if i not in missing_indices]).to(device).long()
+
+    def V(self, vec):
+        temp = vec.clone().reshape(vec.shape[0], -1)
+        out = torch.zeros_like(temp)
+        out[:, self.kept_indices] = temp[:, :self.kept_indices.shape[0]]
+        out[:, self.missing_indices] = temp[:, self.kept_indices.shape[0]:]
+        return out.reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+    def Vt(self, vec):
+        temp = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1).reshape(vec.shape[0], -1)
+        out = torch.zeros_like(temp)
+        out[:, :self.kept_indices.shape[0]] = temp[:, self.kept_indices]
+        out[:, self.kept_indices.shape[0]:] = temp[:, self.missing_indices]
+        return out
+
+    def U(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars
+
+    def add_zeros(self, vec):
+        temp = torch.zeros((vec.shape[0], self.channels * self.img_dim**2), device=vec.device)
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp[:, :reshaped.shape[1]] = reshaped
+        return temp
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+
+        temp = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1).reshape(vec.shape[0], -1)
+        out = torch.zeros_like(temp)
+        out[:, :self.kept_indices.shape[0]] = temp[:, self.kept_indices]
+        out[:, self.kept_indices.shape[0]:] = temp[:, self.missing_indices]
+
+        singulars = self._singulars
+        lambda_t = torch.ones(temp.size(1), device=vec.device)
+        temp_singulars = torch.zeros(temp.size(1), device=vec.device)
+        temp_singulars[:singulars.size(0)] = singulars
+        singulars = temp_singulars
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (
+                    singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, -1)
+        out = out * lambda_t
+
+        result = torch.zeros_like(temp)
+        result[:, self.kept_indices] = out[:, :self.kept_indices.shape[0]]
+        result[:, self.missing_indices] = out[:, self.kept_indices.shape[0]:]
+        return result.reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        temp_vec = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1).reshape(vec.shape[0], -1)
+        out_vec = torch.zeros_like(temp_vec)
+        out_vec[:, :self.kept_indices.shape[0]] = temp_vec[:, self.kept_indices]
+        out_vec[:, self.kept_indices.shape[0]:] = temp_vec[:, self.missing_indices]
+
+        temp_eps = epsilon.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1).reshape(vec.shape[0], -1)
+        out_eps = torch.zeros_like(temp_eps)
+        out_eps[:, :self.kept_indices.shape[0]] = temp_eps[:, self.kept_indices]
+        out_eps[:, self.kept_indices.shape[0]:] = temp_eps[:, self.missing_indices]
+
+        singulars = self._singulars
+        d1_t = torch.ones(temp_vec.size(1), device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(temp_vec.size(1), device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        temp_singulars = torch.zeros(temp_vec.size(1), device=vec.device)
+        temp_singulars[:singulars.size(0)] = singulars
+        singulars = temp_singulars
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(
+                change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, -1)
+        d2_t = d2_t.reshape(1, -1)
+        out_vec = out_vec * d1_t
+        out_eps = out_eps * d2_t
+
+        result_vec = torch.zeros_like(temp_vec)
+        result_vec[:, self.kept_indices] = out_vec[:, :self.kept_indices.shape[0]]
+        result_vec[:, self.missing_indices] = out_vec[:, self.kept_indices.shape[0]:]
+        result_vec = result_vec.reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        result_eps = torch.zeros_like(temp_eps)
+        result_eps[:, self.kept_indices] = out_eps[:, :self.kept_indices.shape[0]]
+        result_eps[:, self.missing_indices] = out_eps[:, self.kept_indices.shape[0]:]
+        result_eps = result_eps.reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        return result_vec + result_eps
+
+#Denoising
+class Denoising(A_functions):
+    def __init__(self, channels, img_dim, device):
+        self._singulars = torch.ones(channels * img_dim**2, device=device)
+
+    def V(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Vt(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars
+
+    def add_zeros(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        if sigma_t < a * sigma_y:
+            factor = (sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y).item()
+            return vec * factor
+        else:
+            return vec
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        if sigma_t >= a * sigma_y:
+            factor = torch.sqrt(sigma_t ** 2 - a ** 2 * sigma_y ** 2).item()
+            return vec * factor
+        else:
+            return vec * sigma_t * eta
+
+#Super Resolution
+class SuperResolution(A_functions):
+    def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
+        assert img_dim % ratio == 0
+        self.img_dim = img_dim
+        self.channels = channels
+        self.y_dim = img_dim // ratio
+        self.ratio = ratio
+        A = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
+        self.U_small, self.singulars_small, self.V_small = torch.svd(A, some=False)
+        self.Vt_small = self.V_small.transpose(0, 1)
+
+    def V(self, vec):
+        #reorder the vector back into patches (because singulars are ordered descendingly)
+        temp = vec.clone().reshape(vec.shape[0], -1)
+        patches = torch.zeros(vec.shape[0], self.channels, self.y_dim**2, self.ratio**2, device=vec.device)
+        patches[:, :, :, 0] = temp[:, :self.channels * self.y_dim**2].view(vec.shape[0], self.channels, -1)
+        for idx in range(self.ratio**2-1):
+            patches[:, :, :, idx+1] = temp[:, (self.channels*self.y_dim**2+idx)::self.ratio**2-1].view(vec.shape[0], self.channels, -1)
+        #multiply each patch by the small V
+        patches = torch.matmul(self.V_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #repatch the patches into an image
+        patches_orig = patches.reshape(vec.shape[0], self.channels, self.y_dim, self.y_dim, self.ratio, self.ratio)
+        recon = patches_orig.permute(0, 1, 2, 4, 3, 5).contiguous()
+        recon = recon.reshape(vec.shape[0], self.channels * self.img_dim ** 2)
+        return recon
+
+    def Vt(self, vec):
+        #extract flattened patches
+        patches = vec.clone().reshape(vec.shape[0], self.channels, self.img_dim, self.img_dim)
+        patches = patches.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        unfold_shape = patches.shape
+        patches = patches.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #multiply each by the small V transposed
+        patches = torch.matmul(self.Vt_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #reorder the vector to have the first entry first (because singulars are ordered descendingly)
+        recon = torch.zeros(vec.shape[0], self.channels * self.img_dim**2, device=vec.device)
+        recon[:, :self.channels * self.y_dim**2] = patches[:, :, :, 0].view(vec.shape[0], self.channels * self.y_dim**2)
+        for idx in range(self.ratio**2-1):
+            recon[:, (self.channels*self.y_dim**2+idx)::self.ratio**2-1] = patches[:, :, :, idx+1].view(vec.shape[0], self.channels * self.y_dim**2)
+        return recon
+
+    def U(self, vec):
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec): #U is 1x1, so U^T = U
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self.singulars_small.repeat(self.channels * self.y_dim**2)
+
+    def add_zeros(self, vec):
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp = torch.zeros((vec.shape[0], reshaped.shape[1] * self.ratio**2), device=vec.device)
+        temp[:, :reshaped.shape[1]] = reshaped
+        return temp
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        singulars = self.singulars_small
+
+        patches = vec.clone().reshape(vec.shape[0], self.channels, self.img_dim, self.img_dim)
+        patches = patches.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches = patches.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        patches = torch.matmul(self.Vt_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+
+        lambda_t = torch.ones(self.ratio ** 2, device=vec.device)
+
+        temp = torch.zeros(self.ratio ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, 1, 1, -1)
+#         print("lambda_t:", lambda_t)
+#         print("V:", self.V_small)
+#         print(lambda_t.size(), self.V_small.size())
+#         print("Sigma_t:", torch.matmul(torch.matmul(self.V_small, torch.diag(lambda_t.reshape(-1))), self.Vt_small))
+        patches = patches * lambda_t
+
+
+        patches = torch.matmul(self.V_small, patches.reshape(-1, self.ratio**2, 1))
+
+        patches = patches.reshape(vec.shape[0], self.channels, self.y_dim, self.y_dim, self.ratio, self.ratio)
+        patches = patches.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches = patches.reshape(vec.shape[0], self.channels * self.img_dim ** 2)
+
+        return patches
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        singulars = self.singulars_small
+
+        patches_vec = vec.clone().reshape(vec.shape[0], self.channels, self.img_dim, self.img_dim)
+        patches_vec = patches_vec.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches_vec = patches_vec.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        patches_eps = epsilon.clone().reshape(vec.shape[0], self.channels, self.img_dim, self.img_dim)
+        patches_eps = patches_eps.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches_eps = patches_eps.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        d1_t = torch.ones(self.ratio ** 2, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.ratio ** 2, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        temp = torch.zeros(self.ratio ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index  * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, 1, 1, -1)
+        d2_t = d2_t.reshape(1, 1, 1, -1)
+        patches_vec = patches_vec * d1_t
+        patches_eps = patches_eps * d2_t
+
+        patches_vec = torch.matmul(self.V_small, patches_vec.reshape(-1, self.ratio**2, 1))
+
+        patches_vec = patches_vec.reshape(vec.shape[0], self.channels, self.y_dim, self.y_dim, self.ratio, self.ratio)
+        patches_vec = patches_vec.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches_vec = patches_vec.reshape(vec.shape[0], self.channels * self.img_dim ** 2)
+
+        patches_eps = torch.matmul(self.V_small, patches_eps.reshape(-1, self.ratio**2, 1))
+
+        patches_eps = patches_eps.reshape(vec.shape[0], self.channels, self.y_dim, self.y_dim, self.ratio, self.ratio)
+        patches_eps = patches_eps.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches_eps = patches_eps.reshape(vec.shape[0], self.channels * self.img_dim ** 2)
+
+        return patches_vec + patches_eps
+
+class SuperResolutionGeneral(SuperResolution):
+    def __init__(self, channels, imgH, imgW, ratio, device): #ratio = 2 or 4
+        assert imgH % ratio == 0 and imgW % ratio == 0
+        self.imgH = imgH
+        self.imgW = imgW
+        self.channels = channels
+        self.yH = imgH // ratio
+        self.yW = imgW // ratio
+        self.ratio = ratio
+        A = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
+        self.U_small, self.singulars_small, self.V_small = torch.svd(A, some=False)
+        self.Vt_small = self.V_small.transpose(0, 1)
+
+    def V(self, vec):
+        #reorder the vector back into patches (because singulars are ordered descendingly)
+        temp = vec.clone().reshape(vec.shape[0], -1)
+        patches = torch.zeros(vec.shape[0], self.channels, self.yH*self.yW, self.ratio**2, device=vec.device)
+        patches[:, :, :, 0] = temp[:, :self.channels * self.yH*self.yW].view(vec.shape[0], self.channels, -1)
+        for idx in range(self.ratio**2-1):
+            patches[:, :, :, idx+1] = temp[:, (self.channels*self.yH*self.yW+idx)::self.ratio**2-1].view(vec.shape[0], self.channels, -1)
+        #multiply each patch by the small V
+        patches = torch.matmul(self.V_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #repatch the patches into an image
+        patches_orig = patches.reshape(vec.shape[0], self.channels, self.yH, self.yW, self.ratio, self.ratio)
+        recon = patches_orig.permute(0, 1, 2, 4, 3, 5).contiguous()
+        recon = recon.reshape(vec.shape[0], self.channels * self.imgH * self.imgW)
+        return recon
+
+    def Vt(self, vec):
+        #extract flattened patches
+        patches = vec.clone().reshape(vec.shape[0], self.channels, self.imgH, self.imgW)
+        patches = patches.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        unfold_shape = patches.shape
+        patches = patches.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #multiply each by the small V transposed
+        patches = torch.matmul(self.Vt_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+        #reorder the vector to have the first entry first (because singulars are ordered descendingly)
+        recon = torch.zeros(vec.shape[0], self.channels * self.imgH*self.imgW, device=vec.device)
+        recon[:, :self.channels * self.yH*self.yW] = patches[:, :, :, 0].view(vec.shape[0], self.channels * self.yH*self.yW)
+        for idx in range(self.ratio**2-1):
+            recon[:, (self.channels*self.yH*self.yW+idx)::self.ratio**2-1] = patches[:, :, :, idx+1].view(vec.shape[0], self.channels * self.yH*self.yW)
+        return recon
+
+    def U(self, vec):
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec): #U is 1x1, so U^T = U
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self.singulars_small.repeat(self.channels * self.yH*self.yW)
+
+    def add_zeros(self, vec):
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp = torch.zeros((vec.shape[0], reshaped.shape[1] * self.ratio**2), device=vec.device)
+        temp[:, :reshaped.shape[1]] = reshaped
+        return temp
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        singulars = self.singulars_small
+
+        patches = vec.clone().reshape(vec.shape[0], self.channels, self.imgH, self.imgW)
+        patches = patches.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches = patches.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        patches = torch.matmul(self.Vt_small, patches.reshape(-1, self.ratio**2, 1)).reshape(vec.shape[0], self.channels, -1, self.ratio**2)
+
+        lambda_t = torch.ones(self.ratio ** 2, device=vec.device)
+
+        temp = torch.zeros(self.ratio ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, 1, 1, -1)
+#         print("lambda_t:", lambda_t)
+#         print("V:", self.V_small)
+#         print(lambda_t.size(), self.V_small.size())
+#         print("Sigma_t:", torch.matmul(torch.matmul(self.V_small, torch.diag(lambda_t.reshape(-1))), self.Vt_small))
+        patches = patches * lambda_t
+
+
+        patches = torch.matmul(self.V_small, patches.reshape(-1, self.ratio**2, 1))
+
+        patches = patches.reshape(vec.shape[0], self.channels, self.yH, self.yW, self.ratio, self.ratio)
+        patches = patches.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches = patches.reshape(vec.shape[0], self.channels * self.imgH * self.imgW)
+
+        return patches
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        singulars = self.singulars_small
+
+        patches_vec = vec.clone().reshape(vec.shape[0], self.channels, self.imgH, self.imgW)
+        patches_vec = patches_vec.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches_vec = patches_vec.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        patches_eps = epsilon.clone().reshape(vec.shape[0], self.channels, self.imgH, self.imgW)
+        patches_eps = patches_eps.unfold(2, self.ratio, self.ratio).unfold(3, self.ratio, self.ratio)
+        patches_eps = patches_eps.contiguous().reshape(vec.shape[0], self.channels, -1, self.ratio ** 2)
+
+        d1_t = torch.ones(self.ratio ** 2, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.ratio ** 2, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        temp = torch.zeros(self.ratio ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index  * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, 1, 1, -1)
+        d2_t = d2_t.reshape(1, 1, 1, -1)
+        patches_vec = patches_vec * d1_t
+        patches_eps = patches_eps * d2_t
+
+        patches_vec = torch.matmul(self.V_small, patches_vec.reshape(-1, self.ratio**2, 1))
+
+        patches_vec = patches_vec.reshape(vec.shape[0], self.channels, self.yH, self.yW, self.ratio, self.ratio)
+        patches_vec = patches_vec.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches_vec = patches_vec.reshape(vec.shape[0], self.channels * self.imgH * self.imgW)
+
+        patches_eps = torch.matmul(self.V_small, patches_eps.reshape(-1, self.ratio**2, 1))
+
+        patches_eps = patches_eps.reshape(vec.shape[0], self.channels, self.yH, self.yW, self.ratio, self.ratio)
+        patches_eps = patches_eps.permute(0, 1, 2, 4, 3, 5).contiguous()
+        patches_eps = patches_eps.reshape(vec.shape[0], self.channels * self.imgH * self.imgW)
+
+        return patches_vec + patches_eps
+
+#Colorization
+class Colorization(A_functions):
+    def __init__(self, img_dim, device):
+        self.channels = 3
+        self.img_dim = img_dim
+        #Do the SVD for the per-pixel matrix
+        A = torch.Tensor([[0.3333, 0.3334, 0.3333]]).to(device)
+        self.U_small, self.singulars_small, self.V_small = torch.svd(A, some=False)
+        self.Vt_small = self.V_small.transpose(0, 1)
+
+    def V(self, vec):
+        #get the needles
+        needles = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1) #shape: B, WA, C'
+        #multiply each needle by the small V
+        needles = torch.matmul(self.V_small, needles.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1, self.channels) #shape: B, WA, C
+        #permute back to vector representation
+        recon = needles.permute(0, 2, 1) #shape: B, C, WA
+        return recon.reshape(vec.shape[0], -1)
+
+    def Vt(self, vec):
+        #get the needles
+        needles = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1) #shape: B, WA, C
+        #multiply each needle by the small V transposed
+        needles = torch.matmul(self.Vt_small, needles.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1, self.channels) #shape: B, WA, C'
+        #reorder the vector so that the first entry of each needle is at the top
+        recon = needles.permute(0, 2, 1).reshape(vec.shape[0], -1)
+        return recon
+
+    def U(self, vec):
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec): #U is 1x1, so U^T = U
+        return self.U_small[0, 0] * vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self.singulars_small.repeat(self.img_dim**2)
+
+    def add_zeros(self, vec):
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp = torch.zeros((vec.shape[0], self.channels * self.img_dim**2), device=vec.device)
+        temp[:, :self.img_dim**2] = reshaped
+        return temp
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        needles = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1)
+
+        needles = torch.matmul(self.Vt_small, needles.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1, self.channels)
+
+        singulars = self.singulars_small
+        lambda_t = torch.ones(self.channels, device=vec.device)
+        temp = torch.zeros(self.channels, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (
+                        singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, 1, self.channels)
+        needles = needles * lambda_t
+
+        needles = torch.matmul(self.V_small, needles.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1, self.channels)
+
+        recon = needles.permute(0, 2, 1).reshape(vec.shape[0], -1)
+        return recon
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        singulars = self.singulars_small
+
+        needles_vec = vec.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1)
+        needles_epsilon = epsilon.clone().reshape(vec.shape[0], self.channels, -1).permute(0, 2, 1)
+
+        d1_t = torch.ones(self.channels, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.channels, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        temp = torch.zeros(self.channels, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(
+                change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, 1, self.channels)
+        d2_t = d2_t.reshape(1, 1, self.channels)
+
+        needles_vec = needles_vec * d1_t
+        needles_epsilon = needles_epsilon * d2_t
+
+        needles_vec = torch.matmul(self.V_small, needles_vec.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1, self.channels)
+        recon_vec = needles_vec.permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        needles_epsilon = torch.matmul(self.V_small, needles_epsilon.reshape(-1, self.channels, 1)).reshape(vec.shape[0], -1,self.channels)
+        recon_epsilon = needles_epsilon.permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        return recon_vec + recon_epsilon
+
+#Walsh-Aadamard Compressive Sensing
+class WalshAadamardCS(A_functions):
+    def fwht(self, vec): #the Fast Walsh Aadamard Transform is the same as its inverse
+        a = vec.reshape(vec.shape[0], self.channels, self.img_dim**2)
+        h = 1
+        while h < self.img_dim**2:
+            a = a.reshape(vec.shape[0], self.channels, -1, h * 2)
+            b = a.clone()
+            a[:, :, :, :h] = b[:, :, :, :h] + b[:, :, :, h:2*h]
+            a[:, :, :, h:2*h] = b[:, :, :, :h] - b[:, :, :, h:2*h]
+            h *= 2
+        a = a.reshape(vec.shape[0], self.channels, self.img_dim**2) / self.img_dim
+        return a
+
+    def __init__(self, channels, img_dim, ratio, perm, device):
+        self.channels = channels
+        self.img_dim = img_dim
+        self.ratio = ratio
+        self.perm = perm
+        self._singulars = torch.ones(channels * img_dim**2 // ratio, device=device)
+
+    def V(self, vec):
+        temp = torch.zeros(vec.shape[0], self.channels, self.img_dim**2, device=vec.device)
+        temp[:, :, self.perm] = vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        return self.fwht(temp).reshape(vec.shape[0], -1)
+
+    def Vt(self, vec):
+        return self.fwht(vec.clone())[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def Ut(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars
+
+    def add_zeros(self, vec):
+        out = torch.zeros(vec.shape[0], self.channels * self.img_dim**2, device=vec.device)
+        out[:, :self.channels * self.img_dim**2 // self.ratio] = vec.clone().reshape(vec.shape[0], -1)
+        return out
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        temp_vec = self.fwht(vec.clone())[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        singulars = self._singulars
+        lambda_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device)
+        temp = torch.zeros(self.channels * self.img_dim ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (
+                    singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, -1)
+        temp_vec = temp_vec * lambda_t
+
+        temp_out = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out[:, :, self.perm] = temp_vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        return self.fwht(temp_out).reshape(vec.shape[0], -1)
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        temp_vec = vec.clone().reshape(
+            vec.shape[0], self.channels, self.img_dim ** 2)[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+        temp_eps = epsilon.clone().reshape(
+            vec.shape[0], self.channels, self.img_dim ** 2)[:, :, self.perm].permute(0, 2, 1).reshape(vec.shape[0], -1)
+
+        d1_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.channels * self.img_dim ** 2, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        singulars = self._singulars
+        temp = torch.zeros(self.channels * self.img_dim ** 2, device=vec.device)
+        temp[:singulars.size(0)] = singulars
+        singulars = temp
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(
+                change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, -1)
+        d2_t = d2_t.reshape(1, -1)
+
+        temp_vec = temp_vec * d1_t
+        temp_eps = temp_eps * d2_t
+
+        temp_out_vec = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out_vec[:, :, self.perm] = temp_vec.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        temp_out_vec = self.fwht(temp_out_vec).reshape(vec.shape[0], -1)
+
+        temp_out_eps = torch.zeros(vec.shape[0], self.channels, self.img_dim ** 2, device=vec.device)
+        temp_out_eps[:, :, self.perm] = temp_eps.clone().reshape(vec.shape[0], -1, self.channels).permute(0, 2, 1)
+        temp_out_eps = self.fwht(temp_out_eps).reshape(vec.shape[0], -1)
+
+        return temp_out_vec + temp_out_eps
+
+#Convolution-based super-resolution
+class SRConv(A_functions):
+    def mat_by_img(self, M, v, dim):
+        return torch.matmul(M, v.reshape(v.shape[0] * self.channels, dim,
+                        dim)).reshape(v.shape[0], self.channels, M.shape[0], dim)
+
+    def img_by_mat(self, v, M, dim):
+        return torch.matmul(v.reshape(v.shape[0] * self.channels, dim,
+                        dim), M).reshape(v.shape[0], self.channels, dim, M.shape[1])
+
+    def __init__(self, kernel, channels, img_dim, device, stride = 1):
+        self.img_dim = img_dim
+        self.channels = channels
+        self.ratio = stride
+        small_dim = img_dim // stride
+        self.small_dim = small_dim
+        #build 1D conv matrix
+        A_small = torch.zeros(small_dim, img_dim, device=device)
+        for i in range(stride//2, img_dim + stride//2, stride):
+            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
+                j_effective = j
+                #reflective padding
+                if j_effective < 0: j_effective = -j_effective-1
+                if j_effective >= img_dim: j_effective = (img_dim - 1) - (j_effective - img_dim)
+                #matrix building
+                A_small[i // stride, j_effective] += kernel[j - i + kernel.shape[0]//2]
+        #get the svd of the 1D conv
+        self.U_small, self.singulars_small, self.V_small = torch.svd(A_small, some=False)
+        ZERO = 3e-2
+        self.singulars_small[self.singulars_small < ZERO] = 0
+        #calculate the singular values of the big matrix
+        self._singulars = torch.matmul(self.singulars_small.reshape(small_dim, 1), self.singulars_small.reshape(1, small_dim)).reshape(small_dim**2)
+        #permutation for matching the singular values. See P_1 in Appendix D.5.
+        self._perm = torch.Tensor([self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim)] + \
+                                  [self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim, self.img_dim)]).to(device).long()
+
+    def V(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.img_dim**2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)[:, :self._perm.shape[0], :]
+        temp[:, self._perm.shape[0]:, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)[:, self._perm.shape[0]:, :]
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by V from the left and by V^T from the right
+        out = self.mat_by_img(self.V_small, temp, self.img_dim)
+        out = self.img_by_mat(out, self.V_small.transpose(0, 1), self.img_dim).reshape(vec.shape[0], -1)
+        return out
+
+    def Vt(self, vec):
+        #multiply the image by V^T from the left and by V from the right
+        temp = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone(), self.img_dim)
+        temp = self.img_by_mat(temp, self.V_small, self.img_dim).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries
+        temp[:, :, :self._perm.shape[0]] = temp[:, :, self._perm]
+        temp = temp.permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.small_dim**2, self.channels, device=vec.device)
+        temp[:, :self.small_dim**2, :] = vec.clone().reshape(vec.shape[0], self.small_dim**2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by U from the left and by U^T from the right
+        out = self.mat_by_img(self.U_small, temp, self.small_dim)
+        out = self.img_by_mat(out, self.U_small.transpose(0, 1), self.small_dim).reshape(vec.shape[0], -1)
+        return out
+
+    def Ut(self, vec):
+        #multiply the image by U^T from the left and by U from the right
+        temp = self.mat_by_img(self.U_small.transpose(0, 1), vec.clone(), self.small_dim)
+        temp = self.img_by_mat(temp, self.U_small, self.small_dim).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries
+        temp = temp.permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars.repeat_interleave(3).reshape(-1)
+
+    def add_zeros(self, vec):
+        reshaped = vec.clone().reshape(vec.shape[0], -1)
+        temp = torch.zeros((vec.shape[0], reshaped.shape[1] * self.ratio**2), device=vec.device)
+        temp[:, :reshaped.shape[1]] = reshaped
+        return temp
+
+#Deblurring
+class Deblurring(A_functions):
+    def mat_by_img(self, M, v):
+        return torch.matmul(M, v.reshape(v.shape[0] * self.channels, self.img_dim,
+                        self.img_dim)).reshape(v.shape[0], self.channels, M.shape[0], self.img_dim)
+
+    def img_by_mat(self, v, M):
+        return torch.matmul(v.reshape(v.shape[0] * self.channels, self.img_dim,
+                        self.img_dim), M).reshape(v.shape[0], self.channels, self.img_dim, M.shape[1])
+
+    def __init__(self, kernel, channels, img_dim, device, ZERO = 3e-2):
+        self.img_dim = img_dim
+        self.channels = channels
+        #build 1D conv matrix
+        A_small = torch.zeros(img_dim, img_dim, device=device)
+        for i in range(img_dim):
+            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
+                if j < 0 or j >= img_dim: continue
+                A_small[i, j] = kernel[j - i + kernel.shape[0]//2]
+        #get the svd of the 1D conv
+        self.U_small, self.singulars_small, self.V_small = torch.svd(A_small, some=False)
+        #ZERO = 3e-2
+        self.singulars_small_orig = self.singulars_small.clone()
+        self.singulars_small[self.singulars_small < ZERO] = 0
+        #calculate the singular values of the big matrix
+        self._singulars_orig = torch.matmul(self.singulars_small_orig.reshape(img_dim, 1), self.singulars_small_orig.reshape(1, img_dim)).reshape(img_dim**2)
+        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim**2)
+        #sort the big matrix singulars and save the permutation
+        self._singulars, self._perm = self._singulars.sort(descending=True) #, stable=True)
+        self._singulars_orig = self._singulars_orig[self._perm]
+
+    def V(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.img_dim**2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by V from the left and by V^T from the right
+        out = self.mat_by_img(self.V_small, temp)
+        out = self.img_by_mat(out, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)
+        return out
+
+    def Vt(self, vec):
+        #multiply the image by V^T from the left and by V from the right
+        temp = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone())
+        temp = self.img_by_mat(temp, self.V_small).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries according to the singular values
+        temp = temp[:, :, self._perm].permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.img_dim**2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by U from the left and by U^T from the right
+        out = self.mat_by_img(self.U_small, temp)
+        out = self.img_by_mat(out, self.U_small.transpose(0, 1)).reshape(vec.shape[0], -1)
+        return out
+
+    def Ut(self, vec):
+        #multiply the image by U^T from the left and by U from the right
+        temp = self.mat_by_img(self.U_small.transpose(0, 1), vec.clone())
+        temp = self.img_by_mat(temp, self.U_small).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries according to the singular values
+        temp = temp[:, :, self._perm].permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars.repeat(1, 3).reshape(-1)
+
+    def add_zeros(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)
+
+    def A_pinv(self, vec):
+        temp = self.Ut(vec)
+        singulars = self._singulars.repeat(1, 3).reshape(-1)
+
+        factors = 1. / singulars
+        factors[singulars == 0] = 0.
+
+        temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] * factors
+        return self.V(self.add_zeros(temp))
+
+    def Lambda(self, vec, a, sigma_y, sigma_t, eta):
+        temp_vec = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone())
+        temp_vec = self.img_by_mat(temp_vec, self.V_small).reshape(vec.shape[0], self.channels, -1)
+        temp_vec = temp_vec[:, :, self._perm].permute(0, 2, 1)
+
+        singulars = self._singulars_orig
+        lambda_t = torch.ones(self.img_dim ** 2, device=vec.device)
+        temp_singulars = torch.zeros(self.img_dim ** 2, device=vec.device)
+        temp_singulars[:singulars.size(0)] = singulars
+        singulars = temp_singulars
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            lambda_t = lambda_t * (-change_index + 1.0) + change_index * (
+                    singulars * sigma_t * (1 - eta ** 2) ** 0.5 / a / sigma_y)
+
+        lambda_t = lambda_t.reshape(1, -1, 1)
+        temp_vec = temp_vec * lambda_t
+
+        temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = temp_vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        out = self.mat_by_img(self.V_small, temp)
+        out = self.img_by_mat(out, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)
+        return out
+
+    def Lambda_noise(self, vec, a, sigma_y, sigma_t, eta, epsilon):
+        temp_vec = vec.clone().reshape(vec.shape[0], self.channels, -1)
+        temp_vec = temp_vec[:, :, self._perm].permute(0, 2, 1)
+
+        temp_eps = epsilon.clone().reshape(vec.shape[0], self.channels, -1)
+        temp_eps = temp_eps[:, :, self._perm].permute(0, 2, 1)
+
+        singulars = self._singulars_orig
+        d1_t = torch.ones(self.img_dim ** 2, device=vec.device) * sigma_t * eta
+        d2_t = torch.ones(self.img_dim ** 2, device=vec.device) * sigma_t * (1 - eta ** 2) ** 0.5
+
+        temp_singulars = torch.zeros(self.img_dim ** 2, device=vec.device)
+        temp_singulars[:singulars.size(0)] = singulars
+        singulars = temp_singulars
+        inverse_singulars = 1. / singulars
+        inverse_singulars[singulars == 0] = 0.
+
+        if a != 0 and sigma_y != 0:
+            change_index = (sigma_t < a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (sigma_t > a * sigma_y * inverse_singulars) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + torch.sqrt(
+                change_index * (sigma_t ** 2 - a ** 2 * sigma_y ** 2 * inverse_singulars ** 2))
+            d2_t = d2_t * (-change_index + 1.0)
+
+            change_index = (singulars == 0) * 1.0
+            d1_t = d1_t * (-change_index + 1.0) + change_index * sigma_t * eta
+            d2_t = d2_t * (-change_index + 1.0) + change_index * sigma_t * (1 - eta ** 2) ** 0.5
+
+        d1_t = d1_t.reshape(1, -1, 1)
+        d2_t = d2_t.reshape(1, -1, 1)
+
+        temp_vec = temp_vec * d1_t
+        temp_eps = temp_eps * d2_t
+
+        temp_vec_new = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)
+        temp_vec_new[:, self._perm, :] = temp_vec
+        out_vec = self.mat_by_img(self.V_small, temp_vec_new.permute(0, 2, 1))
+        out_vec = self.img_by_mat(out_vec, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)
+
+        temp_eps_new = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)
+        temp_eps_new[:, self._perm, :] = temp_eps
+        out_eps = self.mat_by_img(self.V_small, temp_eps_new.permute(0, 2, 1))
+        out_eps = self.img_by_mat(out_eps, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)
+
+        return out_vec + out_eps
+
+#Anisotropic Deblurring
+class Deblurring2D(A_functions):
+    def mat_by_img(self, M, v):
+        return torch.matmul(M, v.reshape(v.shape[0] * self.channels, self.img_dim,
+                        self.img_dim)).reshape(v.shape[0], self.channels, M.shape[0], self.img_dim)
+
+    def img_by_mat(self, v, M):
+        return torch.matmul(v.reshape(v.shape[0] * self.channels, self.img_dim,
+                        self.img_dim), M).reshape(v.shape[0], self.channels, self.img_dim, M.shape[1])
+
+    def __init__(self, kernel1, kernel2, channels, img_dim, device):
+        self.img_dim = img_dim
+        self.channels = channels
+        A_small1 = torch.zeros(img_dim, img_dim, device=device)
+        for i in range(img_dim):
+            for j in range(i - kernel1.shape[0]//2, i + kernel1.shape[0]//2):
+                if j < 0 or j >= img_dim: continue
+                A_small1[i, j] = kernel1[j - i + kernel1.shape[0]//2]
+        A_small2 = torch.zeros(img_dim, img_dim, device=device)
+        for i in range(img_dim):
+            for j in range(i - kernel2.shape[0]//2, i + kernel2.shape[0]//2):
+                if j < 0 or j >= img_dim: continue
+                A_small2[i, j] = kernel2[j - i + kernel2.shape[0]//2]
+        self.U_small1, self.singulars_small1, self.V_small1 = torch.svd(A_small1, some=False)
+        self.U_small2, self.singulars_small2, self.V_small2 = torch.svd(A_small2, some=False)
+        ZERO = 3e-2
+        self.singulars_small1[self.singulars_small1 < ZERO] = 0
+        self.singulars_small2[self.singulars_small2 < ZERO] = 0
+        #calculate the singular values of the big matrix
+        self._singulars = torch.matmul(self.singulars_small1.reshape(img_dim, 1), self.singulars_small2.reshape(1, img_dim)).reshape(img_dim**2)
+        #sort the big matrix singulars and save the permutation
+        self._singulars, self._perm = self._singulars.sort(descending=True) #, stable=True)
+
+    def V(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.img_dim**2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by V from the left and by V^T from the right
+        out = self.mat_by_img(self.V_small1, temp)
+        out = self.img_by_mat(out, self.V_small2.transpose(0, 1)).reshape(vec.shape[0], -1)
+        return out
+
+    def Vt(self, vec):
+        #multiply the image by V^T from the left and by V from the right
+        temp = self.mat_by_img(self.V_small1.transpose(0, 1), vec.clone())
+        temp = self.img_by_mat(temp, self.V_small2).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries according to the singular values
+        temp = temp[:, :, self._perm].permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def U(self, vec):
+        #invert the permutation
+        temp = torch.zeros(vec.shape[0], self.img_dim**2, self.channels, device=vec.device)
+        temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim**2, self.channels)
+        temp = temp.permute(0, 2, 1)
+        #multiply the image by U from the left and by U^T from the right
+        out = self.mat_by_img(self.U_small1, temp)
+        out = self.img_by_mat(out, self.U_small2.transpose(0, 1)).reshape(vec.shape[0], -1)
+        return out
+
+    def Ut(self, vec):
+        #multiply the image by U^T from the left and by U from the right
+        temp = self.mat_by_img(self.U_small1.transpose(0, 1), vec.clone())
+        temp = self.img_by_mat(temp, self.U_small2).reshape(vec.shape[0], self.channels, -1)
+        #permute the entries according to the singular values
+        temp = temp[:, :, self._perm].permute(0, 2, 1)
+        return temp.reshape(vec.shape[0], -1)
+
+    def singulars(self):
+        return self._singulars.repeat(1, 3).reshape(-1)
+
+    def add_zeros(self, vec):
+        return vec.clone().reshape(vec.shape[0], -1)

src/flair/helper_functions.py

@@ -0,0 +1,31 @@
+import os
+import yaml
+import click
+import subprocess
+
+def parse_click_context(ctx):
+    """Parse additional arguments passed via Click context."""
+    extra_args = {}
+    for arg in ctx.args:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+            if key.startswith("--"):
+                key = key[2:]  # Remove leading "--"
+            extra_args[key] = yaml.safe_load(value)
+    return extra_args
+
+def generate_captions_with_seesr(pseudo_inv_dir, output_caption_file):
+    """Generate captions using the SEESR model."""
+    command = [
+        "conda",
+        "run",
+        "-n",
+        "seesr",
+        "python",
+        "/home/erbachj/scratch2/projects/var_post_samp/scripts/generate_caption.py",
+        "--input_dir",
+        pseudo_inv_dir,
+        "--output_file",
+        output_caption_file,  # Corrected argument name
+    ]
+    subprocess.run(command, check=True)

src/flair/pipelines/model_loader.py

@@ -0,0 +1,97 @@
+# functions to load models from config
+import numpy as np
+import torch
+import re
+import os
+import math
+from diffusers import (
+    BitsAndBytesConfig
+)
+from diffusers import AutoencoderTiny
+
+
+from flair.pipelines import sd3
+
+
+
+@torch.no_grad()
+def load_sd3(config, device):
+    if isinstance(device, list):
+        device = device[0]
+    if config["quantize"]:
+        nf4_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+    else:
+        nf4_config = None
+    if config["model"] == "SD3.5-large":
+        pipe = sd3.SD3Wrapper.from_pretrained(
+            "stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16, quantization_config=nf4_config
+        )
+    elif config["model"] == "SD3.5-large-turbo":
+        pipe = sd3.SD3Wrapper.from_pretrained(
+            "stabilityai/stable-diffusion-3.5-large-turbo", torch_dtype=torch.bfloat16, quantization_config=nf4_config,
+        )
+    else:
+        pipe = sd3.SD3Wrapper.from_pretrained("stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch.float16, quantization_config=nf4_config)
+    # maybe use tiny autoencoder
+    if config["use_tiny_ae"]:
+        pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd3", torch_dtype=torch.float16)
+
+    # encode prompts
+    inp_kwargs_list = []
+    prompts = config["prompt"]
+    pipe._guidance_scale = config["guidance"]
+    pipe._joint_attention_kwargs = {"ip_adapter_image_embeds": None}
+    for prompt in prompts:
+        print(f"Generating prompt embeddings for: {prompt}")
+        pipe.text_encoder.to(device).to(torch.bfloat16)
+        pipe.text_encoder_2.to(device).to(torch.bfloat16)
+        pipe.text_encoder_3.to(device).to(torch.bfloat16)
+        # encode
+        (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+        ) = pipe.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt,
+            prompt_3=prompt,
+            negative_prompt=config["negative_prompt"],
+            negative_prompt_2=config["negative_prompt"],
+            negative_prompt_3=config["negative_prompt"],
+            do_classifier_free_guidance=pipe.do_classifier_free_guidance,
+            prompt_embeds=None,
+            negative_prompt_embeds=None,
+            pooled_prompt_embeds=None,
+            negative_pooled_prompt_embeds=None,
+            device=device,
+            clip_skip=None,
+            num_images_per_prompt=1,
+            max_sequence_length=256,
+            lora_scale=None,
+        )
+
+        if pipe.do_classifier_free_guidance:
+            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)
+        inp_kwargs = {
+            "prompt_embeds": prompt_embeds,
+            "pooled_prompt_embeds": pooled_prompt_embeds,
+            "guidance": config["guidance"],
+        }
+        inp_kwargs_list.append(inp_kwargs)
+    pipe.vae.to(device).to(torch.bfloat16)
+    pipe.transformer.to(device).to(torch.bfloat16)
+
+
+    return pipe, inp_kwargs_list
+
+def load_model(config, device=["cuda"]):
+    if re.match(r"SD3*", config["model"]):
+        return load_sd3(config, device)
+    else:
+        raise ValueError(f"Unknown model type {config['model']}")

src/flair/pipelines/sd3.py

@@ -0,0 +1,111 @@
+import torch
+from typing import Dict, Any
+from diffusers.pipelines.stable_diffusion_3 import pipeline_stable_diffusion_3
+from flair.pipelines import utils
+import tqdm
+
+class SD3Wrapper(pipeline_stable_diffusion_3.StableDiffusion3Pipeline):
+    def to(self, device, kwargs):
+        self.transformer.to(device)
+        self.vae.to(device)
+        return self
+
+    def get_timesteps(self, n_steps, device, ts_min=0):
+        # Create a linear schedule for timesteps
+        timesteps = torch.linspace(1, ts_min, n_steps+2, device=device, dtype=torch.float32)
+        return timesteps[1:-1]  # Exclude the first and last timesteps
+
+    def single_step(
+        self,
+        img_latent: torch.Tensor,
+        t: torch.Tensor,
+        kwargs: Dict[str, Any],
+        is_noised_latent = False,
+    ):  
+        if "noise" in kwargs:
+            noise = kwargs["noise"].detach()
+            alpha = kwargs["inv_alpha"]
+            if alpha == "tsqrt":
+                alpha = t**0.5  # * 0.75
+            elif alpha == "t":
+                alpha = t
+            elif alpha == "sine":
+                alpha = torch.sin(t * 3.141592653589793/2)
+            elif alpha == "1-t":
+                alpha = 1 - t
+            elif alpha == "1-t*0.5":
+                alpha = (1 - t)*0.5
+            elif alpha == "1-t*0.9":
+                alpha = (1 - t)*0.9
+            elif alpha == "t**1/3":
+                alpha = t**(1/3)
+            elif alpha == "(1-t)**0.5":
+                alpha = (1-t)**0.5
+            elif alpha == "((1-t)*0.8)**0.5":
+                alpha = (1-t*0.8)**0.5
+            elif alpha == "(1-t)**2":
+                alpha = (1-t)**2
+            # alpha = t * kwargs["inv_alpha"]
+            noise = (alpha) * noise + (1-alpha**2)**0.5 * torch.randn_like(img_latent)
+            # noise = noise / noise.std()
+            # noise = noise / (1- 2*alpha*(1-alpha))**0.5
+            # noise = noise + alpha * torch.randn_like(img_latent)
+        else:
+            noise = torch.randn_like(img_latent)
+        if is_noised_latent:
+            noised_latent = img_latent
+        else:
+            noised_latent = t * noise + (1 - t) * img_latent
+        latent_model_input = torch.cat([noised_latent] * 2) if self.do_classifier_free_guidance else noised_latent
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timestep = t.expand(latent_model_input.shape[0])
+        noise_pred = self.transformer(
+            hidden_states=latent_model_input.to(img_latent.dtype),
+            timestep=(timestep*1000).to(img_latent.dtype),
+            encoder_hidden_states=kwargs["prompt_embeds"].repeat(img_latent.shape[0], 1, 1),
+            pooled_projections=kwargs["pooled_prompt_embeds"].repeat(img_latent.shape[0], 1),
+            joint_attention_kwargs=None,
+            return_dict=False,
+        )[0]
+        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)
+  
+        eps = utils.v_to_eps(noise_pred, t, noised_latent)
+        return eps, noise, (1 - t), t, noise_pred
+    
+    def encode(self, img):
+        # Encode the image into latent space
+        
+        img_latent = self.vae.encode(img, return_dict=False)[0]
+        if hasattr(img_latent, "sample"):
+            img_latent = img_latent.sample()
+        img_latent = (img_latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        return img_latent
+
+    def decode(self, img_latent):
+        # Decode the latent representation back to image space
+        img = self.vae.decode(img_latent / self.vae.config.scaling_factor + self.vae.config.shift_factor, return_dict=False)[0]
+        return img
+
+    def denoise(self, pseudo_inv, kwargs, inverse=False):
+        # get timesteps
+        timesteps = torch.linspace(1, 0, kwargs["n_steps"], device=pseudo_inv.device, dtype=pseudo_inv.dtype)
+        sigmas = timesteps
+        if inverse:
+            timesteps = timesteps.flip(0)
+            sigmas = sigmas.flip(0)
+        
+        # make a single step
+        for i, t in tqdm.tqdm(enumerate(timesteps[:-1]), desc="Denoising", total=len(timesteps)-1):
+            eps, noise, _, t, v = self.single_step(
+                pseudo_inv,
+                t.to("cuda")*1000,
+                kwargs,
+                is_noised_latent=True,
+            )
+            # step
+            sigma_next = sigmas[i+1]
+            sigma_t = sigmas[i]
+            pseudo_inv = pseudo_inv + v * (sigma_next - sigma_t)
+        return pseudo_inv

src/flair/pipelines/utils.py

@@ -0,0 +1,114 @@
+import torch
+import numpy as np
+
+
+def eps_from_v(z_0, z_t, sigma_t):
+    return (z_t - z_0) / sigma_t
+
+def v_to_eps(v, t, x_t):
+    """
+    function to compute the epsilon parametrization from the velocity field
+    with x_t = t * x_0 + (1 - t) * x_1 with x_0 ~ N(0,I)
+    """
+    eps_t = (1-t)*v + x_t
+    return eps_t
+
+
+def clip_gradients(gradients, clip_value):
+    grad_norm = gradients.norm(dim=2)
+    mask = grad_norm > clip_value
+    mask_exp = mask[:, :, None].expand_as(gradients)
+    gradients[mask_exp] = (
+        gradients[mask_exp]
+        / grad_norm[:, :, None].expand_as(gradients)[mask_exp]
+        * clip_value
+    )
+    return gradients
+
+
+class Adam:
+
+    def __init__(self, parameters, lr=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
+        self.lr = lr
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.epsilon = epsilon
+        self.t = 0
+        self.m = torch.zeros_like(parameters)
+        self.v = torch.zeros_like(parameters)
+
+    def step(self, params, grad) -> torch.Tensor:
+        self.t += 1
+
+        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
+        self.v = self.beta2 * self.v + (1 - self.beta2) * grad**2
+
+        m_hat = self.m / (1 - self.beta1**self.t)
+        v_hat = self.v / (1 - self.beta2**self.t)
+
+        # check if self.lr is callable
+        if callable(self.lr):
+            lr = self.lr(self.t - 1)
+        else:
+            lr = self.lr
+        update = lr * m_hat / (torch.sqrt(v_hat) + self.epsilon)
+
+        return params - update
+
+
+def make_cosine_decay_schedule(
+    init_value: float,
+    total_steps: int,
+    alpha: float = 0.0,
+    exponent: float = 1.0,
+    warmup_steps=0,
+):
+    def schedule(count):
+        if count < warmup_steps:
+            # linear up
+            return (init_value / warmup_steps) * count
+        else:
+            # half cosine down
+            decay_steps = total_steps - warmup_steps
+            count = min(count - warmup_steps, decay_steps)
+            cosine_decay = 0.5 * (1 + np.cos(np.pi * count / decay_steps))
+            decayed = (1 - alpha) * cosine_decay**exponent + alpha
+            return init_value * decayed
+
+    return schedule
+
+
+def make_linear_decay_schedule(
+    init_value: float, total_steps: int, final_value: float = 0, warmup_steps=0
+):
+    def schedule(count):
+        if count < warmup_steps:
+            # linear up
+            return (init_value / warmup_steps) * count
+        else:
+            # linear down
+            decay_steps = total_steps - warmup_steps
+            count = min(count - warmup_steps, decay_steps)
+            return init_value - (init_value - final_value) * count / decay_steps
+
+    return schedule
+
+
+def clip_norm_(tensor, max_norm):
+    norm = tensor.norm()
+    if norm > max_norm:
+        tensor.mul_(max_norm / norm)
+
+
+def lr_warmup(step, warmup_steps):
+    return min(1.0, step / max(warmup_steps, 1))
+
+
+def linear_decay_lambda(step, warmup_steps, decay_steps, total_steps):
+    if step < warmup_steps:
+        min(1.0, step / max(warmup_steps, 1))
+    else:
+        # linear down
+        # decay_steps = total_steps - warmup_steps
+        count = min(step - warmup_steps, decay_steps)
+        return 1 - count / decay_steps

src/flair/scheduling.py

@@ -0,0 +1,882 @@
+# Copyright 2023 Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import (
+    KarrasDiffusionSchedulers,
+    SchedulerMixin,
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete
+class EulerDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    pred_original_sample: Optional[torch.FloatTensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.FloatTensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.FloatTensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear` or `scaled_linear`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        interpolation_type(`str`, defaults to `"linear"`, *optional*):
+            The interpolation type to compute intermediate sigmas for the scheduler denoising steps. Should be on of
+            `"linear"` or `"log_linear"`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps. You can use a combination of `offset=1` and
+            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
+            Diffusion.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        prediction_type: str = "epsilon",
+        interpolation_type: str = "linear",
+        use_karras_sigmas: Optional[bool] = False,
+        sigma_min: Optional[float] = None,
+        sigma_max: Optional[float] = None,
+        timestep_spacing: str = "linspace",
+        timestep_type: str = "discrete",  # can be "discrete" or "continuous"
+        steps_offset: int = 0,
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(
+                beta_start, beta_end, num_train_timesteps, dtype=torch.float32
+            )
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = (
+                torch.linspace(
+                    beta_start**0.5,
+                    beta_end**0.5,
+                    num_train_timesteps,
+                    dtype=torch.float32,
+                )
+                ** 2
+            )
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(
+                f"{beta_schedule} does is not implemented for {self.__class__}"
+            )
+
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        if rescale_betas_zero_snr:
+            # Close to 0 without being 0 so first sigma is not inf
+            # FP16 smallest positive subnormal works well here
+            self.alphas_cumprod[-1] = 2**-24
+
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        timesteps = np.linspace(
+            0, num_train_timesteps - 1, num_train_timesteps, dtype=float
+        )[::-1].copy()
+
+        sigmas = torch.from_numpy(sigmas[::-1].copy()).to(dtype=torch.float32)
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        # setable values
+        self.num_inference_steps = None
+
+        # TODO: Support the full EDM scalings for all prediction types and timestep types
+        if timestep_type == "continuous" and prediction_type == "v_prediction":
+            self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas])
+        else:
+            self.timesteps = timesteps
+
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self.is_scale_input_called = False
+        self.use_karras_sigmas = use_karras_sigmas
+
+        self._step_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigmas_for_inference = False
+
+    @property
+    def init_noise_sigma(self):
+        # standard deviation of the initial noise distribution
+        max_sigma = (
+            max(self.sigmas) if isinstance(self.sigmas, list) else self.sigmas.max()
+        )
+        if self.config.timestep_spacing in ["linspace", "trailing"]:
+            return max_sigma
+
+        return (max_sigma**2 + 1) ** 0.5
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increae 1 after each scheduler step.
+        """
+        return self._step_index
+
+    def scale_model_input(
+        self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor]
+    ) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        sigma = self.sigmas[self.step_index]
+        sample = sample / ((sigma**2 + 1) ** 0.5)
+
+        self.is_scale_input_called = True
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        timesteps=None,
+        device: Union[str, torch.device] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+        self.num_inference_steps = num_inference_steps
+
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+        if self.config.timestep_spacing == "linspace":
+            timesteps = np.linspace(
+                0,
+                self.config.num_train_timesteps - 1,
+                num_inference_steps,
+                dtype=np.float32,
+            )[::-1].copy()
+        elif self.config.timestep_spacing == "leading":
+            step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (
+                (np.arange(0, num_inference_steps) * step_ratio)
+                .round()[::-1]
+                .copy()
+                .astype(np.float32)
+            )
+            timesteps += self.config.steps_offset
+        elif self.config.timestep_spacing == "trailing":
+            step_ratio = self.config.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (
+                (np.arange(self.config.num_train_timesteps, 0, -step_ratio))
+                .round()
+                .copy()
+                .astype(np.float32)
+            )
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+            )
+
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        log_sigmas = np.log(sigmas)
+
+        if self.config.interpolation_type == "linear":
+            sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+        elif self.config.interpolation_type == "log_linear":
+            sigmas = (
+                torch.linspace(
+                    np.log(sigmas[-1]), np.log(sigmas[0]), num_inference_steps + 1
+                )
+                .exp()
+                .numpy()
+            )
+        else:
+            raise ValueError(
+                f"{self.config.interpolation_type} is not implemented. Please specify interpolation_type to either"
+                " 'linear' or 'log_linear'"
+            )
+
+        if self.use_karras_sigmas:
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_steps=self.num_inference_steps
+            )
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]
+            )
+
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+
+        # TODO: Support the full EDM scalings for all prediction types and timestep types
+        if (
+            self.config.timestep_type == "continuous"
+            and self.config.prediction_type == "v_prediction"
+        ):
+            self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas]).to(
+                device=device
+            )
+        else:
+            self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(
+                device=device
+            )
+
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+        self._step_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigmas_for_inference = True
+
+    def _sigma_to_t(self, sigma, log_sigmas):
+        # get log sigma
+        log_sigma = np.log(np.maximum(sigma, 1e-10))
+
+        # get distribution
+        dists = log_sigma - log_sigmas[:, np.newaxis]
+
+        # get sigmas range
+        low_idx = (
+            np.cumsum((dists >= 0), axis=0)
+            .argmax(axis=0)
+            .clip(max=log_sigmas.shape[0] - 2)
+        )
+        high_idx = low_idx + 1
+
+        low = log_sigmas[low_idx]
+        high = log_sigmas[high_idx]
+
+        # interpolate sigmas
+        w = (low - log_sigma) / (low - high)
+        w = np.clip(w, 0, 1)
+
+        # transform interpolation to time range
+        t = (1 - w) * low_idx + w * high_idx
+        t = t.reshape(sigma.shape)
+        return t
+
+    # Copied from https://github.com/crowsonkb/k-diffusion/blob/686dbad0f39640ea25c8a8c6a6e56bb40eacefa2/k_diffusion/sampling.py#L17
+    def _convert_to_karras(
+        self, in_sigmas: torch.FloatTensor, num_inference_steps
+    ) -> torch.FloatTensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    def derivative_sigma_t(self, step):
+        rho = 7.0
+        sigma_max = self.sigmas[0]
+        sigma_min = self.sigmas[-1]
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        derivative_sigma_t = (
+            -7
+            * (max_inv_rho + (step) / 1000 * (min_inv_rho - max_inv_rho)) ** 6
+            * (min_inv_rho - max_inv_rho)
+        )
+        return derivative_sigma_t
+
+    def _init_step_index(self, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.to(self.timesteps.device)
+
+        index_candidates = (self.timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        if len(index_candidates) > 1:
+            step_index = index_candidates[1]
+        else:
+            step_index = index_candidates[0]
+
+        self._step_index = step_index.item()
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+        likelihood_grad=None,
+        likelihood_weight=None,
+        regularizer_weight=None,
+        step_index: Optional[int] = None,
+    ) -> Union[EulerDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+                tuple.
+
+        Returns:
+            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if (
+            isinstance(timestep, int)
+            or isinstance(timestep, torch.IntTensor)
+            or isinstance(timestep, torch.LongTensor)
+        ):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if not self.is_scale_input_called:
+            logger.warning(
+                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
+                "See `StableDiffusionPipeline` for a usage example."
+            )
+
+        # if self.step_index is None:
+        #     self._init_step_index(timestep)
+        if step_index is None:
+            if self.step_index is None:
+                self._init_step_index(timestep)
+            step_index = torch.tensor(self.step_index).unsqueeze(0)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if likelihood_grad is not None:
+            likelihood_grad = likelihood_grad.to(torch.float32)
+
+        sigma = self.sigmas[step_index][:, None, None, None].to(model_output.device)
+
+        gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
+        condition = torch.logical_and(s_tmin <= sigma, sigma <= s_tmax)
+        gamma = torch.where(condition, gamma, torch.tensor(0))
+
+        noise = randn_tensor(
+            model_output.shape,
+            dtype=model_output.dtype,
+            device=model_output.device,
+            generator=generator,
+        )
+
+        eps = noise * s_noise
+        sigma_hat = sigma * (gamma + 1)
+
+        # if gamma > 0:
+        sample = torch.where(
+            gamma > 0, sample + eps * (sigma_hat**2 - sigma**2) ** 0.5, sample
+        )
+        # sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
+
+        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+        # NOTE: "original_sample" should not be an expected prediction_type but is left in for
+        # backwards compatibility
+        if (
+            self.config.prediction_type == "original_sample"
+            or self.config.prediction_type == "sample"
+        ):
+            pred_original_sample = model_output
+        elif self.config.prediction_type == "epsilon":
+            pred_original_sample = sample - sigma_hat * model_output
+        elif self.config.prediction_type == "v_prediction":
+            # denoised = model_output * c_out + input * c_skip
+            pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (
+                sample / (sigma**2 + 1)
+            )
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+            )
+
+        # 2. Convert to an ODE derivative
+        derivative = (sample - pred_original_sample) / sigma_hat
+        if likelihood_grad is not None:
+            likelihood_grad = likelihood_grad * (
+                torch.norm(derivative) / torch.norm(likelihood_grad)
+            )
+            derivative = (
+                regularizer_weight * derivative + likelihood_weight * likelihood_grad
+            )
+
+        dt = (
+            self.sigmas[step_index + 1].to(sigma_hat.device)[:, None, None, None]
+            - sigma_hat
+        )
+
+        prev_sample = sample + derivative * dt
+
+        # Cast sample back to model compatible dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        if self._step_index is not None:
+            self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return EulerDiscreteSchedulerOutput(
+            prev_sample=prev_sample, pred_original_sample=pred_original_sample
+        )
+
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype
+        )
+        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32
+            )
+            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        noisy_samples = original_samples + noise * sigma
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+
+
+class HeunDiscreteScheduler(EulerDiscreteScheduler):
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        timesteps=None,
+        device: Union[str, torch.device] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+        self.num_inference_steps = num_inference_steps
+
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+        if self.config.timestep_spacing == "linspace":
+            timesteps = np.linspace(
+                0,
+                self.config.num_train_timesteps - 1,
+                num_inference_steps,
+                dtype=np.float32,
+            )[::-1].copy()
+        elif self.config.timestep_spacing == "leading":
+            step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (
+                (np.arange(0, num_inference_steps) * step_ratio)
+                .round()[::-1]
+                .copy()
+                .astype(np.float32)
+            )
+            timesteps += self.config.steps_offset
+        elif self.config.timestep_spacing == "trailing":
+            step_ratio = self.config.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (
+                (np.arange(self.config.num_train_timesteps, 0, -step_ratio))
+                .round()
+                .copy()
+                .astype(np.float32)
+            )
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+            )
+
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        log_sigmas = np.log(sigmas)
+
+        if self.config.interpolation_type == "linear":
+            sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+        elif self.config.interpolation_type == "log_linear":
+            sigmas = (
+                torch.linspace(
+                    np.log(sigmas[-1]), np.log(sigmas[0]), num_inference_steps + 1
+                )
+                .exp()
+                .numpy()
+            )
+        else:
+            raise ValueError(
+                f"{self.config.interpolation_type} is not implemented. Please specify interpolation_type to either"
+                " 'linear' or 'log_linear'"
+            )
+
+        if self.use_karras_sigmas:
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_steps=self.num_inference_steps
+            )
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]
+            )
+
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+
+        sigmas = torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2)])
+
+        # TODO: Support the full EDM scalings for all prediction types and timestep types
+        if (
+            self.config.timestep_type == "continuous"
+            and self.config.prediction_type == "v_prediction"
+        ):
+            self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas]).to(
+                device=device
+            )
+        else:
+            self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(
+                device=device
+            )
+
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+        self._step_index = None
+        self.sigmas = self.sigmas.to("cpu")
+
+        # empty dt and derivative
+        self.prev_derivative = None
+        self.dt = None
+
+    def step(
+        self,
+        model_output: Union[torch.Tensor, np.ndarray],
+        timestep: Union[float, torch.Tensor],
+        sample: Union[torch.Tensor, np.ndarray],
+        return_dict: bool = True,
+    ) -> Union[EulerDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or tuple.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        if self.state_in_first_order:
+            sigma = self.sigmas[self.step_index]
+            sigma_next = self.sigmas[self.step_index + 1]
+        else:
+            # 2nd order / Heun's method
+            sigma = self.sigmas[self.step_index - 1]
+            sigma_next = self.sigmas[self.step_index]
+
+        # currently only gamma=0 is supported. This usually works best anyways.
+        # We can support gamma in the future but then need to scale the timestep before
+        # passing it to the model which requires a change in API
+        gamma = 0
+        sigma_hat = sigma * (gamma + 1)  # Note: sigma_hat == sigma for now
+
+        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+        if self.config.prediction_type == "epsilon":
+            sigma_input = sigma_hat if self.state_in_first_order else sigma_next
+            pred_original_sample = sample - sigma_input * model_output
+        elif self.config.prediction_type == "v_prediction":
+            sigma_input = sigma_hat if self.state_in_first_order else sigma_next
+            pred_original_sample = model_output * (
+                -sigma_input / (sigma_input**2 + 1) ** 0.5
+            ) + (sample / (sigma_input**2 + 1))
+        elif self.config.prediction_type == "sample":
+            pred_original_sample = model_output
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+            )
+
+        if self.config.clip_sample:
+            pred_original_sample = pred_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        if self.state_in_first_order:
+            # 2. Convert to an ODE derivative for 1st order
+            derivative = (sample - pred_original_sample) / sigma_hat
+            # 3. delta timestep
+            dt = sigma_next - sigma_hat
+
+            # store for 2nd order step
+            self.prev_derivative = derivative
+            self.dt = dt
+            self.sample = sample
+        else:
+            # 2. 2nd order / Heun's method
+            derivative = (sample - pred_original_sample) / sigma_next
+            derivative = (self.prev_derivative + derivative) / 2
+
+            # 3. take prev timestep & sample
+            dt = self.dt
+            sample = self.sample
+
+            # free dt and derivative
+            # Note, this puts the scheduler in "first order mode"
+            self.prev_derivative = None
+            self.dt = None
+            self.sample = None
+
+        prev_sample = sample + derivative * dt
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return EulerDiscreteSchedulerOutput(
+            prev_sample=prev_sample, pred_original_sample=pred_original_sample
+        )
+
+    @property
+    def state_in_first_order(self):
+        return self.dt is None
+
+    @property
+    def order(self):
+        return 2

src/flair/utils.py

@@ -0,0 +1,406 @@
+import random
+from pathlib import Path
+import os
+import torch
+from torchmetrics.functional.image import (
+    peak_signal_noise_ratio,
+    learned_perceptual_image_patch_similarity,
+)
+from PIL import Image
+from skimage.color import rgb2lab, lab2rgb
+import numpy as np
+import cv2
+from torchvision import transforms
+
+
+RESAMPLE_MODE = Image.BICUBIC
+
+
+def skip_iterator(iterator, skip):
+    for i, item in enumerate(iterator):
+        if i % skip == 0:
+            yield item
+
+
+def generate_output_structure(output_dir, subfolders=[]):
+    """
+    Generate a directory structure for the output. and return the paths to the subfolders. as template
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(exist_ok=True, parents=True)
+    output_paths = []
+    for subfolder in subfolders:
+        output_paths.append(Path(os.path.join(output_dir, subfolder)))
+        (output_paths[-1]).mkdir(exist_ok=True, parents=True)
+        output_paths[-1] = os.path.join(output_paths[-1], "{}.png")
+    return output_paths
+
+
+def find_files(path, ext="png"):
+    if os.path.isdir(path):
+        path = Path(path)
+        sorted_files = sorted(list(path.glob(f"*.{ext}")))
+        return sorted_files
+    else:
+        return [path]
+
+
+def load_guidance_image(path, size=None):
+    """
+    Load an image and convert it to a tensor.
+    Args: path to the image
+    returns: tensor of the image of shape (1, 3, H, W)
+    """
+    img = Image.open(path)
+    img = img.convert("RGB")
+    tf = transforms.Compose([
+    transforms.Resize(size),
+    transforms.CenterCrop(size),
+    transforms.ToTensor()
+    ])
+    img = tf(img) * 2 - 1
+    return img.unsqueeze(0)
+
+
+def yield_images(path, ext="png", size=None):
+    files = find_files(path, ext)
+    for file in files:
+        yield load_guidance_image(file, size)
+
+
+def yield_videos(paths, ext="png", H=None, W=None, n_frames=61):
+    for path in paths:
+        yield read_video(path, H, W, n_frames)
+
+
+def read_video(path, H=None, W=None, n_frames=61) -> list[Image]:
+    path = Path(path)
+    frames = []
+
+    if Path(path).is_dir():
+        files = sorted(list(path.glob("*.png")))
+        for file in files[:n_frames]:
+            image = Image.open(file)
+            image.load()
+            if H is not None and W is not None:
+                image = image.resize((W, H), resample=Image.BICUBIC)
+            # to tensor
+            image = (
+                torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1)
+                / 255.0
+                * 2
+                - 1
+            )
+            frames.append(image)
+        H, W = frames[0].size()[-2:]
+        frames = torch.stack(frames).unsqueeze(0)
+        return frames, (10, H, W)
+
+    capture = cv2.VideoCapture(str(path))
+    fps = capture.get(cv2.CAP_PROP_FPS)
+
+    while True:
+        success, frame = capture.read()
+        if not success:
+            break
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame = Image.fromarray(frame)
+        if H is not None and W is not None:
+            frame = frame.resize((W, H), resample=Image.BICUBIC)
+        # to tensor
+        frame = (
+            torch.tensor(np.array(frame), dtype=torch.float32).permute(2, 0, 1)
+            / 255.0
+            * 2
+            - 1
+        )
+        frames.append(frame)
+
+    capture.release()
+    # to torch
+    frames = torch.stack(frames).unsqueeze(0)
+    return frames, (fps, W, H)
+
+
+def resize_video(
+    video: list[Image], width, height, resample_mode=RESAMPLE_MODE
+) -> list[Image]:
+    frames_lr = []
+    for frame in video:
+        frame_lr = frame.resize((width, height), resample_mode)
+        frames_lr.append(frame_lr)
+    return frames_lr
+
+
+def export_to_video(
+    video_frames,
+    output_video_path=None,
+    fps=8,
+    put_numbers=False,
+    annotations=None,
+    fourcc="mp4v",
+):
+    fourcc = cv2.VideoWriter_fourcc(*fourcc)  # codec
+    writer = cv2.VideoWriter(output_video_path, fourcc, fps, video_frames[0].size)
+    for i, frame in enumerate(video_frames):
+        frame = cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR)
+
+        if put_numbers:
+            text_position = (frame.shape[1] - 60, 30)
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 1
+            font_color = (255, 255, 255)
+            line_type = 2
+            cv2.putText(
+                frame,
+                f"{i + 1}",
+                text_position,
+                font,
+                font_scale,
+                font_color,
+                line_type,
+            )
+
+        if annotations:
+            annotation = annotations[i]
+            frame = draw_bodypose(
+                frame, annotation["candidates"], annotation["subsets"]
+            )
+
+        writer.write(frame)
+    writer.release()
+
+
+def export_images(frames: list[Image], dir_name):
+    dir_name = Path(dir_name)
+    dir_name.mkdir(exist_ok=True, parents=True)
+    for i, frame in enumerate(frames):
+        frame.save(dir_name / f"{i:05d}.png")
+
+
+def vid2tensor(images: list[Image]) -> torch.Tensor:
+    # PIL to numpy
+    if not isinstance(images, list):
+        raise ValueError()
+    images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+    images = np.stack(images, axis=0)
+
+    if images.ndim == 3:
+        # L mode, add luminance channel
+        images = np.expand_dims(images, -1)
+
+    # numpy to torch
+    images = torch.from_numpy(images.transpose(0, 3, 1, 2))
+    return images
+
+
+def compute_metrics(
+    source: list[Image],
+    output: list[Image],
+    output_lq: list[Image],
+    target: list[Image],
+) -> dict:
+    psnr_ab = torch.tensor(
+        np.mean(compute_color_metrics(output, target)["psnr_ab"])
+    ).float()
+
+    source = vid2tensor(source)
+    output = vid2tensor(output)
+    output_lq = vid2tensor(output_lq)
+    target = vid2tensor(target)
+
+    mse = ((output - target) ** 2).mean()
+    psnr = peak_signal_noise_ratio(output, target, data_range=1.0, dim=(1, 2, 3))
+    # lpips = learned_perceptual_image_patch_similarity(output, target)
+
+    mse_source = ((output_lq - source) ** 2).mean()
+    psnr_source = peak_signal_noise_ratio(
+        output_lq, source, data_range=1.0, dim=(1, 2, 3)
+    )
+
+    return {
+        "mse": mse.detach().cpu().item(),
+        "psnr": psnr.detach().cpu().item(),
+        "psnr_ab": psnr_ab.detach().cpu().item(),
+        "mse_source": mse_source.detach().cpu().item(),
+        "psnr_source": psnr_source.detach().cpu().item(),
+    }
+
+
+def compute_psnr_ab(x, y_gt, pp_max=202.33542248):
+    """Computes the PSNR of the ab color channels.
+
+    Note that the CIE-Lab space is asymmetric.
+    The maximum size for the 2 channels of the ab subspace is approximately 202.3354...
+    pp_max: Approximated maximum swing for the ab channels of the CIE-Lab color space
+        max_{x \in CIE-Lab} {x_a x_b} - min_{x \in CIE-Lab} {x_a x_b}
+    """
+    assert (
+        len(x.shape) == 3
+    ), f"Expecting data of the size HW2 but found {x.shape}; This should be a,b channels of CIE-Lab Space"
+    assert (
+        len(y_gt.shape) == 3
+    ), f"Expecting data of the size HW2 but found {y_gt.shape}; This should be a,b channels of CIE-Lab Space"
+    assert (
+        x.shape == y_gt.shape
+    ), f"Expecting data to have identical shape but found {y_gt.shape} != {x.shape}"
+
+    H, W, C = x.shape
+    assert (
+        C == 2
+    ), f"This function assumes that both x & y are both the ab channels of the CIE-Lab Space"
+
+    MSE = np.sum((x - y_gt) ** 2) / (H * W * C)  # C=2, two channels
+    MSE = np.clip(MSE, 1e-12, np.inf)
+
+    PSNR_ab = 10 * np.log10(pp_max**2) - 10 * np.log10(MSE)
+
+    return PSNR_ab
+
+
+def compute_color_metrics(out: list[Image], target: list[Image]):
+    if len(out) != len(target):
+        raise ValueError("Videos do not have same length")
+
+    metrics = {"psnr_ab": []}
+
+    for out_frame, target_frame in zip(out, target):
+        out_frame, target_frame = np.asarray(out_frame), np.asarray(target_frame)
+        out_frame_lab, target_frame_lab = rgb2lab(out_frame), rgb2lab(target_frame)
+
+        psnr_ab = compute_psnr_ab(out_frame_lab[..., 1:3], target_frame_lab[..., 1:3])
+        metrics["psnr_ab"].append(psnr_ab.item())
+
+    return metrics
+
+
+def to_device(sample, device):
+    result = {}
+    for key, val in sample.items():
+        if isinstance(val, torch.Tensor):
+            result[key] = val.to(device)
+        elif isinstance(val, list):
+            new_val = []
+            for e in val:
+                if isinstance(e, torch.Tensor):
+                    new_val.append(e.to(device))
+                else:
+                    new_val.append(val)
+            result[key] = new_val
+        else:
+            result[key] = val
+    return result
+
+
+def seed_all(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def adapt_unet(unet, lora_rank=None, in_conv_mode="zeros"):
+    # adapt conv_in
+    kernel = unet.conv_in.weight.data
+    if in_conv_mode == "zeros":
+        new_kernel = torch.zeros(320, 4, 3, 3, dtype=kernel.dtype, device=kernel.device)
+    elif in_conv_mode == "reflect":
+        new_kernel = kernel[:, 4:].clone()
+    else:
+        raise NotImplementedError
+    unet.conv_in.weight.data = torch.cat([kernel, new_kernel], dim=1)
+    if in_conv_mode == "reflect":
+        unet.conv_in.weight.data *= 2.0 / 3.0
+    unet.conv_in.in_channels = 12
+
+    if lora_rank is not None:
+        from peft import LoraConfig
+
+        types = (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Linear, torch.nn.Embedding)
+        target_modules = [
+            (n, m) for n, m in unet.named_modules() if isinstance(m, types)
+        ]
+        # identify parameters (not modules) that will not be lora'd
+        for _, m in target_modules:
+            m.requires_grad_(False)
+        not_adapted = [p for p in unet.parameters() if p.requires_grad]
+
+        unet_lora_config = LoraConfig(
+            r=lora_rank,
+            lora_alpha=lora_rank,
+            init_lora_weights="gaussian",
+            # target_modules=["to_k", "to_q", "to_v", "to_out.0"],
+            target_modules=[n for n, _ in target_modules],
+        )
+        # the following line sets all parameters except the loras to non-trainable
+        unet.add_adapter(unet_lora_config)
+
+        unet.conv_in.requires_grad_()
+        for p in not_adapted:
+            p.requires_grad_()
+
+
+def repeat_infinite(iterable):
+    def repeated():
+        while True:
+            yield from iterable
+
+    return repeated
+
+
+class CPUAdam:
+    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        self.params = list(params)
+        self.lr = lr
+        self.betas = betas
+        self.eps = eps
+        self.weight_decay = weight_decay
+        # keep this in main memory to save VRAM
+        self.state = {
+            param: {
+                "step": 0,
+                "exp_avg": torch.zeros_like(param, device="cpu"),
+                "exp_avg_sq": torch.zeros_like(param, device="cpu"),
+            }
+            for param in self.params
+        }
+
+    def step(self):
+        for param in self.params:
+            if param.grad is None:
+                continue
+
+            grad = param.grad.data.cpu()
+            if self.weight_decay != 0:
+                grad.add_(param.data, alpha=self.weight_decay)
+
+            state = self.state[param]
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            beta1, beta2 = self.betas
+
+            state["step"] += 1
+
+            exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+            denom = exp_avg_sq.sqrt().add_(self.eps)
+
+            step_size = (
+                self.lr
+                * (1 - beta2 ** state["step"]) ** 0.5
+                / (1 - beta1 ** state["step"])
+            )
+
+            # param.data.add_((-step_size * (exp_avg / denom)).cuda())
+
+            param.data.addcdiv_(exp_avg.cuda(), denom.cuda(), value=-step_size)
+
+    def zero_grad(self):
+        for param in self.params:
+            param.grad = None
+
+    def state_dict(self):
+        return self.state
+
+    def set_lr(self, lr):
+        self.lr = lr

src/flair/utils/blur_util.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+import numpy as np
+import scipy
+
+from flair.utils.motionblur import Kernel as MotionKernel
+
+class Blurkernel(nn.Module):
+    def __init__(self, blur_type='gaussian', kernel_size=31, std=3.0, device=None):
+        super().__init__()
+        self.blur_type = blur_type
+        self.kernel_size = kernel_size
+        self.std = std
+        self.device = device
+        self.seq = nn.Sequential(
+            nn.ReflectionPad2d(self.kernel_size//2),
+            nn.Conv2d(3, 3, self.kernel_size, stride=1, padding=0, bias=False, groups=3)
+        )
+
+        self.weights_init()
+
+    def forward(self, x):
+        return self.seq(x)
+
+    def weights_init(self):
+        if self.blur_type == "gaussian":
+            n = np.zeros((self.kernel_size, self.kernel_size))
+            n[self.kernel_size // 2,self.kernel_size // 2] = 1
+            k = scipy.ndimage.gaussian_filter(n, sigma=self.std)
+            k = torch.from_numpy(k)
+            self.k = k
+            for name, f in self.named_parameters():
+                f.data.copy_(k)
+        elif self.blur_type == "motion":
+            k = MotionKernel(size=(self.kernel_size, self.kernel_size), intensity=self.std).kernelMatrix
+            k = torch.from_numpy(k)
+            self.k = k
+            for name, f in self.named_parameters():
+                f.data.copy_(k)
+
+    def update_weights(self, k):
+        if not torch.is_tensor(k):
+            k = torch.from_numpy(k).to(self.device)
+        for name, f in self.named_parameters():
+            f.data.copy_(k)
+
+    def get_kernel(self):
+        return self.k