first

Dockerfile

@@ -0,0 +1,19 @@
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 
+
+RUN apt-get update && apt-get install -y \
+    python3-pip \
+    libgl1-mesa-glx \
+    libglib2.0-0 python3-dev libglib2.0-0 libgl1-mesa-glx \
+    cmake tmux git curl wget gcc build-essential \
+    && rm -rf /var/lib/apt/lists/* 
+
+RUN pip install --no-cache-dir --upgrade pip
+RUN pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121
+ADD requirements.txt requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+WORKDIR /workspace
+
+
+
+
+

DockerfileOcto

@@ -0,0 +1,31 @@
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 
+
+RUN apt-get update && apt-get install -y \
+    python3-pip \
+    libgl1-mesa-glx \
+    libglib2.0-0 python3-dev libglib2.0-0 libgl1-mesa-glx \
+    cmake tmux git curl wget gcc build-essential \
+    && rm -rf /var/lib/apt/lists/* 
+
+WORKDIR /app
+RUN pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121
+ADD requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+ADD guided_diffusion/ /app/guided_diffusion
+ADD dataset/ /app/dataset
+ADD networks/ /app/networks
+ADD utils/ /app/utils
+ADD models /app/models
+ADD server.py /app/server.py
+
+ADD custommodel.py /app/custommodel.py
+
+
+ARG SERVING_PORT=8000
+ENV SERVING_PORT=$SERVING_PORT
+EXPOSE $SERVING_PORT
+
+CMD python3 -m server
+
+

README.md

@@ -0,0 +1,129 @@
+# Distil-DIRE
+Distil-DIRE is a lightweight version of DIRE, which can be used for real-time applications. Instead of calculating DIRE image directly, Distl-DIRE aims to reconstruct the features of corresponding DIRE image forwared by a image-net pretrained classifier with one-step noise of DDIM inversion. ([Paper Link](https://arxiv.org/abs/2406.00856))
+![overview](distil.png)
+
+### Pretrained DistilDIRE Checkpoints
+| Dataset | Model | Link |
+| --- | --- | --- |
+| CelebA-HQ | 224x224 | [link]() |
+| ImageNet | 224x224 | [link]() |
+
+
+
+
+## Pretrained ADM diffusion model
+We use image-net pretrained unconditional ADM diffusion model for feature reconstruction. You can download the pretrained model from the following link:
+https://openaipublic.blob.core.windows.net/diffusion/jul-2021/512x512_diffusion.pt
+
+or you can use the following script to download the model:
+```bash
+wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/512x512_diffusion.pt -O models/512x512-adm.pt
+```
+
+## Data Preparation
+Before training the model on your own dataset, you need to prepare the dataset in the following format:
+```bash
+mydataset/train|val|test
+└── images
+    ├── fakes
+    │   └──img1.png...
+    ├── reals
+        └──rimg1.png...
+```
+
+After preparing the dataset, you can calculate the epsilons and dire images for the dataset using the following script:
+```bash
+bash compute_dire_eps.sh
+```
+
+After running the script, you will have the following directory structure:
+```bash
+mydataset/train|val|test
+└── images
+    ├── fakes
+    │   └──img1.png...
+    ├── reals
+        └──rimg1.png...
+└── eps
+    ├── fakes
+    │   └──img1.pt...
+    ├── reals
+        └──rimg1.pt...
+└── dire
+    ├── fakes
+    │   └──img1.png...
+    ├── reals
+        └──rimg1.png...
+``` 
+For eps and dire calculation we set the DDIM steps to 20. This should be same when inference.
+
+### Train
+For training Distil-DIRE, be sure to have `datasets` directory in the root of the project and your dataset inside the `datasets` directory. 
+```
+torchrun --standalone --nproc_per_node 8 -m train --batch 128 --exp_name truemedia-distil-dire-mydataset --datasets mydataset --datasets_test mytestset --epoch 100 --lr 1e-4
+
+```
+
+#### Fine-tuning
+You can also fine-tune the model on your own dataset. For fine-tuning, you need to provide the path to the pretrained model. 
+```bash
+torchrun --standalone --nproc_per_node 8 -m train --batch 128 --exp_name truemedia-distil-dire-mydataset --datasets mydataset --datasets_test mytestset --epoch 100 --lr 1e-4 --pretrained_weights YOUR_PRETRAINED_MODEL_PATH
+```
+ 
+
+### Test
+For testing the model, you can use the following script:
+```bash
+python3 -m test --test True --datasets mydataset --pretrained_weights YOUR_PRETRAINED_MODEL_PATH
+```
+
+
+### with Docker 
+```
+export DOCKER_REGISTRY="YOUR_NAME" # Put your Docker Hub username here  
+export DATE=`date +%Y%m%d` # Get the current date
+
+# Build the Docker image for development
+docker build -t "$DOCKER_REGISTRY/distil-dire:dev-$DATE" -f Dockerfile .
+
+
+# Push your docker image to docker hub
+docker login
+docker push "$DOCKER_REGISTRY/distil-dire:dev-$DATE"
+
+```
+
+
+# Devl env 
+```
+export WORKDIR="YOUR_WORKDIR" # Put your working directory here
+docker run --gpus=all --name=truemedia_gpu_all_distildire -v "$WORKDIR:/workspace/" -ti -e  "$DOCKER_REGISTRY/distil-dire:dev-$DATE"
+
+# work inside the container (/workspace)
+```
+
+### Note
+* This repo runs on ADM diffusion model (256x256, unconditional) trained on ImageNet 1k dataset and ResNet-50 classifier trained on ImageNet 1k dataset. 
+* Minimum requirements: 1 GPU, 10GB VRAM
+
+
+## Licenses
+This work is licensed under the Creative Commons Attribution-NonCommercial 4.0 license.
+You are free to read, share, and modify this code as long as you keep the original author attribution and non-commercial license.
+Please see [this site](https://creativecommons.org/licenses/by-nc/4.0/) for detailed legal terms.
+
+## Acknowledgments
+Our code is developed based on [DIRE](https://github.com/ZhendongWang6/DIRE), [guided-diffusion](https://github.com/openai/guided-diffusion) and [CNNDetection](https://github.com/peterwang512/CNNDetection). Thanks for their sharing codes and models.
+
+## Citation
+If you find this work useful for your research, please cite our paper:
+```
+@misc{lim2024distildire,
+      title={DistilDIRE: A Small, Fast, Cheap and Lightweight Diffusion Synthesized Deepfake Detection}, 
+      author={Yewon Lim and Changyeon Lee and Aerin Kim and Oren Etzioni},
+      year={2024},
+      eprint={2406.00856},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

check_compress.py

@@ -0,0 +1,42 @@
+from PIL import Image
+import piexif
+
+def check_image_compression(image_path):
+    try:
+        with Image.open(image_path) as img:
+            format = img.format
+            # Common compressed formats
+            if format in ["JPEG", "JPG", "PNG", "GIF", "WEBP"]:
+                print(f"The image format is {format}, which is typically compressed.")
+            else:
+                print(f"The image format is {format}, which is less likely to be compressed.")
+            
+            if format in ["JPEG", "JPG"]:
+                print(img.info)
+                exif_dict = piexif.load(img.info.get('exif', ''))
+                if "0th" in exif_dict and piexif.ImageIFD.Compression in exif_dict["0th"]:
+                    compression = exif_dict["0th"][piexif.ImageIFD.Compression]
+                    if compression != 1:  # 1 means uncompressed
+                        print(f"The JPEG image is compressed. Compression value: {compression}")
+                    else:
+                        print("The JPEG image is not compressed.")
+                else:
+                    print("No EXIF compression metadata found for JPEG image.")
+            elif format == "PNG":
+                if img.info.get("interlace"):
+                    print("The PNG image is interlaced (which can be a form of compression).")
+                else:
+                    print("The PNG image is not interlaced.")
+            elif format == "WEBP":
+                if img.info.get("lossless", False):
+                    print("The WebP image is lossless (not compressed).")
+                else:
+                    print("The WebP image is compressed.")
+            else:
+                print(f"No additional compression information available for format: {format}")
+    except Exception as e:
+        print(f"Error: {e}")
+
+# Example usage
+image_path = '/home/ubuntu/y1/DistilDIRE/datasets/truemedia-political/images/fakes/ZPNOX270fe2dJ6fHvG1mVaHwHM0.jpg'  # Replace with your image path
+check_image_compression(image_path)

compute_dire_eps.sh

@@ -0,0 +1,15 @@
+## set MODEL_PATH, num_samples, has_subfolder, images_dir, recons_dir, dire_dir
+DATA_ROOT=("/home/ubuntu/y1/DistilDIRE/datasets/truemedia-total")
+SAVE_ROOT=("/home/ubuntu/y1/DistilDIRE/datasets/truemedia-total")
+
+MODEL_PATH="models/256x256-adm.pt" # imagenet pretrained adm (unconditional, 512x512)
+SAMPLE_FLAGS="--batch_size 32" # ddim20 is forced
+PREPROCESS_FLAGS="--compute_dire True --compute_eps True"
+
+for i in 0 
+do
+    SAVE_FLAGS="--data_root ${DATA_ROOT[$i]} --save_root ${SAVE_ROOT[$i]}"
+    echo "Running on ${DATA_ROOT[$i]} with save root ${SAVE_ROOT[$i]}"
+    torchrun --standalone --nproc_per_node 1 -m guided_diffusion.compute_dire_eps --model_path $MODEL_PATH $PREPROCESS_FLAGS $SAMPLE_FLAGS $SAVE_FLAGS
+done
+# torchrun --standalone --nproc_per_node 8 -m train --batch 128 --exp_name tm-global-scale --datasets y1-global-truemedia --epoch 40 --lr 1e-4

custommodel.py

@@ -0,0 +1,173 @@
+import argparse
+import json
+from time import perf_counter
+from datetime import datetime
+import torchvision.transforms as transforms
+from torchvision.utils import save_image
+import torchvision.transforms.functional as TF
+from torchvision.transforms import Compose, Resize, CenterCrop
+from torchvision.io import decode_jpeg, encode_jpeg
+
+
+from PIL import Image 
+import os
+import torch
+from utils.sanic_utils import *
+import typing
+import requests
+import time  # Import the time module
+from guided_diffusion.compute_dire_eps import dire_get_first_step_noise, create_argparser
+from networks.distill_model import DistilDIREOnlyEPS, DistilDIRE
+from guided_diffusion.guided_diffusion.script_util import (
+    create_model_and_diffusion,
+    model_and_diffusion_defaults,
+    dict_parse
+)
+
+
+def download_file(input_path):
+    """
+    Download a file from a given URL and save it locally if input_path is a URL.
+    If input_path is a local file path and the file exists, skip the download.
+
+    :param input_path: The URL of the file to download or a local file path.
+    :return: The local filepath to the downloaded or existing file.
+    """
+    # Check if input_path is a URL
+    if input_path.startswith(('http://', 'https://')):
+        # Extract filename from the URL
+        # Splits the URL by '/' and get the last part
+        filename = input_path.split('/')[-1]
+
+        # Ensure the filename does not contain query parameters if present in the URL
+        # Splits the filename by '?' and get the first part
+        filename = filename.split('?')[0]
+
+        # put jpg extension if not present
+        if '.' not in filename:
+            filename += ".jpg"
+
+        # Define the local path where the file will be saved
+        local_filepath = os.path.join('.', filename)
+
+        # Check if file already exists locally
+        if os.path.isfile(local_filepath):
+            print(f"The file already exists locally: {local_filepath}")
+            return local_filepath
+
+        # Start timing the download
+        start_time = time.time()
+
+        # Send a GET request to the URL
+        response = requests.get(input_path, stream=True)
+
+        # Raise an exception if the request was unsuccessful
+        response.raise_for_status()
+
+        # Open the local file in write-binary mode
+        with open(local_filepath, 'wb') as file:
+            # Write the content of the response to the local file
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)
+
+        # End timing the download
+        end_time = time.time()
+
+        # Calculate the download duration
+        download_duration = end_time - start_time
+
+        print(
+            f"Downloaded file saved to {local_filepath} in {download_duration:.2f} seconds.")
+
+    else:
+        # Assume input_path is a local file path
+        local_filepath = input_path
+        # Check if the specified local file exists
+        if not os.path.isfile(local_filepath):
+            raise FileNotFoundError(f"No such file: '{local_filepath}'")
+        print(f"Using existing file: {local_filepath}")
+
+    return local_filepath
+
+
+class CustomModel:
+    """Wrapper for a DIRE model."""
+
+    def __init__(self, net='DIRE', num_frames=15, ckpt='truemedia-global-scaled.pth'):
+        self.net = net
+        self.num_frames = num_frames
+        
+        # self.model =  DistilDIREOnlyEPS('cuda').to('cuda')
+        self.model = DistilDIRE('cuda').to('cuda')
+        self.trans = transforms.Compose((transforms.Resize(256, antialias=True), transforms.CenterCrop((256, 256)),))
+        
+        self._load_state_dict(ckpt)
+        
+        args = create_argparser()
+        args['timestep_respacing'] = 'ddim20'
+        adm_model, diffusion = create_model_and_diffusion(**dict_parse(args, model_and_diffusion_defaults().keys()))
+        adm_model.load_state_dict(torch.load(args['model_path'], map_location="cpu"))
+        adm_model.cuda()
+        adm_model.convert_to_fp16()
+        adm_model.eval()
+        self.adm_model = adm_model
+        self.diffusion = diffusion
+        self.args = args    
+
+
+    def _load_state_dict(self, ckpt):
+        print(f"Loading the model from {ckpt}...")
+        state_dict = torch.load(ckpt, map_location="cpu")['model'] 
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+        self.model.cuda()
+        print("The model is successfully loaded")
+
+
+    def _forward_dire_img(self, img_path, save_dire=True, thr=0.5):
+        img = Image.open(img_path).convert("RGB")
+        img = TF.to_tensor(img)
+        img = self.trans(img).cuda() * 2 - 1
+        img = img.unsqueeze(0)
+        with torch.no_grad():
+            eps = dire_get_first_step_noise(img, self.adm_model, self.diffusion, self.args, "cuda")
+            # eps = eps.detach().cpu()
+            # ext = img_path.split('.')[-1]
+            # eps_path = img_path.replace(f".{ext}", ".pt")
+            # torch.save(eps, eps_path)
+
+            # eps = torch.load(eps_path, weights_only=True, mmap=True).cuda()
+            # os.remove(eps_path)
+            prob = self.model(img, eps)['logit'].sigmoid()
+            
+        return {"df_probability": prob.median().item(), "prediction": real_or_fake_thres(prob.median().item(), thr)}
+
+
+    def predict(self, inputs: typing.Dict[str, str]) -> typing.Dict[str, str]:
+        file_path = inputs.get('file_path', None)
+        video_file = download_file(file_path)
+
+        if os.path.isfile(video_file):
+            try:
+                if is_image(video_file):
+                    print(f"{self.net} is being run.")
+                    return self._forward_dire_img(video_file)
+                
+                else:
+                    print(
+                        f"Invalid media file: {video_file}. Please provide a valid video/img file.")
+                    return {"msg":  f"Invalid media file: {video_file}. Please provide a valid video/img file."}
+            except Exception as e:
+                print(f"An error occurred: {str(e)}")
+                return {"msg": f"An error occurred: {str(e)}"}
+        else:
+            print(f"The file {video_file} does not exist.")
+            return {"msg": f"The file {video_file} does not exist."}
+        
+
+    @classmethod
+    def fetch(cls) -> None:
+        cls()
+

dataset/__init__.py

@@ -0,0 +1 @@
+from .truemedia_dataset import TMDistilDireDataset, TMDireDataset, TMEPSOnlyDataset, TMIMGOnlyDataset

dataset/test.ipynb

@@ -0,0 +1,85 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "172d80aab4da42c8b6a723cd4fc43f6d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/625 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 11\u001b[0m\n\u001b[1;32m      8\u001b[0m dataloader \u001b[38;5;241m=\u001b[39m DataLoader(dataset, batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m128\u001b[39m, shuffle\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, num_workers\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4\u001b[39m)\n\u001b[1;32m     10\u001b[0m st \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m tqdm(dataloader):\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[1;32m     14\u001b[0m end \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/tqdm/notebook.py:250\u001b[0m, in \u001b[0;36mtqdm_notebook.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    249\u001b[0m     it \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m(tqdm_notebook, \u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__iter__\u001b[39m()\n\u001b[0;32m--> 250\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m it:\n\u001b[1;32m    251\u001b[0m         \u001b[38;5;66;03m# return super(tqdm...) will not catch exception\u001b[39;00m\n\u001b[1;32m    252\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m    253\u001b[0m \u001b[38;5;66;03m# NB: except ... [ as ...] breaks IPython async KeyboardInterrupt\u001b[39;00m\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/tqdm/std.py:1181\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1178\u001b[0m time \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_time\n\u001b[1;32m   1180\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1181\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m iterable:\n\u001b[1;32m   1182\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m   1183\u001b[0m         \u001b[38;5;66;03m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   1184\u001b[0m         \u001b[38;5;66;03m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:631\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    628\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    629\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    630\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 631\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    632\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    633\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    634\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    635\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:1329\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1326\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_data(data)\n\u001b[1;32m   1328\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_shutdown \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m-> 1329\u001b[0m idx, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1330\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m   1331\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable:\n\u001b[1;32m   1332\u001b[0m     \u001b[38;5;66;03m# Check for _IterableDatasetStopIteration\u001b[39;00m\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:1295\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._get_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1291\u001b[0m     \u001b[38;5;66;03m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001b[39;00m\n\u001b[1;32m   1292\u001b[0m     \u001b[38;5;66;03m# need to call `.task_done()` because we don't use `.join()`.\u001b[39;00m\n\u001b[1;32m   1293\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1294\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m-> 1295\u001b[0m         success, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_try_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1296\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m success:\n\u001b[1;32m   1297\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
+      "File \u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:1133\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1120\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_try_get_data\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout\u001b[38;5;241m=\u001b[39m_utils\u001b[38;5;241m.\u001b[39mMP_STATUS_CHECK_INTERVAL):\n\u001b[1;32m   1121\u001b[0m     \u001b[38;5;66;03m# Tries to fetch data from `self._data_queue` once for a given timeout.\u001b[39;00m\n\u001b[1;32m   1122\u001b[0m     \u001b[38;5;66;03m# This can also be used as inner loop of fetching without timeout, with\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1130\u001b[0m     \u001b[38;5;66;03m# Returns a 2-tuple:\u001b[39;00m\n\u001b[1;32m   1131\u001b[0m     \u001b[38;5;66;03m#   (bool: whether successfully get data, any: data if successful else None)\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1133\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_data_queue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1134\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m (\u001b[38;5;28;01mTrue\u001b[39;00m, data)\n\u001b[1;32m   1135\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1136\u001b[0m         \u001b[38;5;66;03m# At timeout and error, we manually check whether any worker has\u001b[39;00m\n\u001b[1;32m   1137\u001b[0m         \u001b[38;5;66;03m# failed. Note that this is the only mechanism for Windows to detect\u001b[39;00m\n\u001b[1;32m   1138\u001b[0m         \u001b[38;5;66;03m# worker failures.\u001b[39;00m\n",
+      "File \u001b[0;32m/usr/lib/python3.10/multiprocessing/queues.py:113\u001b[0m, in \u001b[0;36mQueue.get\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m block:\n\u001b[1;32m    112\u001b[0m     timeout \u001b[38;5;241m=\u001b[39m deadline \u001b[38;5;241m-\u001b[39m time\u001b[38;5;241m.\u001b[39mmonotonic()\n\u001b[0;32m--> 113\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_poll\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m    114\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m Empty\n\u001b[1;32m    115\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_poll():\n",
+      "File \u001b[0;32m/usr/lib/python3.10/multiprocessing/connection.py:257\u001b[0m, in \u001b[0;36m_ConnectionBase.poll\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    255\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_closed()\n\u001b[1;32m    256\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_readable()\n\u001b[0;32m--> 257\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_poll\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/usr/lib/python3.10/multiprocessing/connection.py:424\u001b[0m, in \u001b[0;36mConnection._poll\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    423\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_poll\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout):\n\u001b[0;32m--> 424\u001b[0m     r \u001b[38;5;241m=\u001b[39m \u001b[43mwait\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    425\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mbool\u001b[39m(r)\n",
+      "File \u001b[0;32m/usr/lib/python3.10/multiprocessing/connection.py:931\u001b[0m, in \u001b[0;36mwait\u001b[0;34m(object_list, timeout)\u001b[0m\n\u001b[1;32m    928\u001b[0m     deadline \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mmonotonic() \u001b[38;5;241m+\u001b[39m timeout\n\u001b[1;32m    930\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 931\u001b[0m     ready \u001b[38;5;241m=\u001b[39m \u001b[43mselector\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mselect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    932\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ready:\n\u001b[1;32m    933\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m [key\u001b[38;5;241m.\u001b[39mfileobj \u001b[38;5;28;01mfor\u001b[39;00m (key, events) \u001b[38;5;129;01min\u001b[39;00m ready]\n",
+      "File \u001b[0;32m/usr/lib/python3.10/selectors.py:416\u001b[0m, in \u001b[0;36m_PollLikeSelector.select\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    414\u001b[0m ready \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m    415\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 416\u001b[0m     fd_event_list \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_selector\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpoll\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    417\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mInterruptedError\u001b[39;00m:\n\u001b[1;32m    418\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m ready\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "from truemedia_dataset import TMDistilDireDataset\n",
+    "from torch.utils.data import DataLoader\n",
+    "import time \n",
+    "from tqdm.auto import tqdm \n",
+    "\n",
+    "\n",
+    "dataset= TMDistilDireDataset(\"/workspace/DistilDIRE/datasets/distil-train-imagenet\")\n",
+    "dataloader = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=4)\n",
+    "\n",
+    "st = time.time()\n",
+    "for _ in tqdm(dataloader):\n",
+    "    pass\n",
+    "\n",
+    "end = time.time()\n",
+    "print(f\"Time taken: {end-st:.2f} seconds\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

dataset/truemedia_dataset.py

@@ -0,0 +1,147 @@
+from torch.utils.data import Dataset
+import torch
+import torchvision.transforms.functional as TF
+from torchvision.transforms import Compose, Resize, CenterCrop
+from torchvision.io import decode_jpeg, encode_jpeg
+from glob import glob
+import os.path as osp
+from PIL import Image
+import random
+import os 
+TARGET_COMP = 0.1
+
+class TMDistilDireDataset(Dataset):
+    def __init__(self, root, prepared_dire=True):
+        self.root = root
+        self.__fake_img_paths = [p for p in glob(osp.join(root, 'images/fakes/', '*')) if p.split('.')[-1].lower() in ['jpg', 'jpeg', 'png', 'webp']]
+        self.__real_img_paths = [p for p in glob(osp.join(root, 'images/reals/', '*')) if p.split('.')[-1].lower() in ['jpg', 'jpeg', 'png', 'webp']]
+        self.prepared_dire = prepared_dire
+        self.transform = Compose([Resize(256, antialias='True'), CenterCrop((256, 256))])
+        
+
+        # (imgs, dire, eps, isfake)
+        if prepared_dire:
+            self.fake_paths = list(map(lambda x: (x, x.replace('/images/', '/dire/'), x.replace('/images/', '/eps/').split('.')[0]+'.pt', True), self.__fake_img_paths))
+            self.real_paths = list(map(lambda x: (x, x.replace('/images/', '/dire/'), x.replace('/images/', '/eps/').split('.')[0]+'.pt', False), self.__real_img_paths))
+        else:
+            self.fake_paths = list(map(lambda x: (x, "", "", True), self.__fake_img_paths))
+            self.real_paths = list(map(lambda x: (x, "", "", False), self.__real_img_paths))
+        random.shuffle(self.fake_paths)
+        random.shuffle(self.real_paths)
+        self.img_paths = self.fake_paths + self.real_paths
+        # if len(self.real_paths) > len(self.fake_paths) * 3:
+        #     self.img_paths = self.fake_paths + self.real_paths[:len(self.fake_paths) * 3]
+            
+        img_paths = []
+        for img_path, dire_path, eps_path, isfake in self.img_paths:
+            try:
+                Image.open(img_path)
+                img_paths.append((img_path, dire_path, eps_path, isfake))
+            except:
+                continue
+        self.img_paths = img_paths
+
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, idx):
+        img_path, dire_path, eps_path, isfake = self.img_paths[idx]
+        img = Image.open(img_path).convert('RGB')
+        # w, h = img.size
+        # fsize = os.stat(img_path).st_size
+
+        # img = (TF.to_tensor(img)*255).to(torch.uint8)
+
+        # comp = fsize/(w*h)
+        # comp_quality = min(TARGET_COMP/comp * 100, 100)
+        # comp_quality = max(comp_quality, 1)
+        # img = decode_jpeg(encode_jpeg(img, quality=int(comp_quality))) 
+        # img = img / 255.
+        img = TF.to_tensor(img)*2 - 1
+
+        
+        if self.prepared_dire:
+            img = self.transform(img)
+            dire = Image.open(dire_path).convert('RGB')
+            dire = TF.to_tensor(dire)*2 - 1
+            dire = self.transform(dire)
+            eps = torch.load(eps_path, weights_only=True, mmap=True)
+            # eps = torch.zeros_like(img)
+            assert img.shape[1:] == dire.shape[1:] == eps.shape[1:], f"Shape mismatch: {img.shape[1:]}, {dire.shape[1:]}, {eps.shape[1:]}"
+            
+        else:
+            img = self.transform(img)
+            dire = torch.zeros_like(img)
+            eps = torch.zeros_like(img)
+
+        return (img, dire, eps, isfake), (img_path, dire_path, eps_path)
+
+
+
+class TMIMGOnlyDataset(TMDistilDireDataset):
+    def __init__(self, root, istrain=True):
+        super().__init__(root, prepared_dire=True)
+        self.istrain=istrain
+
+    def __getitem__(self, idx):
+       
+        img_path, dire_path, eps_path, isfake = self.img_paths[idx]
+        img = Image.open(img_path).convert('RGB')
+        img = self.transform(img)
+        img = TF.to_tensor(img)*2 - 1
+        eps = torch.zeros_like(img)
+        dire = torch.zeros_like(img)
+    
+        if torch.rand(1) < 0.3 and self.istrain:
+            img = TF.hflip(img)
+        return (img, dire, eps, isfake), (img_path, dire_path, eps_path)
+        
+  
+
+
+class TMEPSOnlyDataset(TMDistilDireDataset):
+    def __init__(self, root, istrain=True):
+        super().__init__(root, prepared_dire=True)
+        img_paths = []
+        for img_path, dire_path, eps_path, isfake in self.img_paths:
+            if not osp.exists(eps_path) or not osp.exists(img_path):
+                # print(f"File not found: {eps_path} or {img_path}")
+                continue 
+            try:
+                eps = torch.load(eps_path, weights_only=True, mmap=True)
+                img_paths.append((img_path, dire_path, eps_path, isfake))
+            except Exception as e:
+                print(e)
+                continue
+        self.img_paths = img_paths
+        self.istrain=istrain
+
+    def __getitem__(self, idx):
+       
+        img_path, dire_path, eps_path, isfake = self.img_paths[idx]
+        img = Image.open(img_path).convert('RGB')
+        img = TF.to_tensor(img)*2 - 1
+        img = self.transform(img)
+        eps = torch.load(eps_path, weights_only=True, mmap=True)
+        dire = torch.zeros_like(img)
+    
+        if torch.rand(1) < 0.3 and self.istrain:
+            img = TF.hflip(img)
+            eps = TF.hflip(eps)
+        
+        return (img, dire, eps, isfake), (img_path, dire_path, eps_path)
+        
+    
+# This is for reproducing DIRE results
+class TMDireDataset(TMDistilDireDataset):
+    def __init__(self, root):
+        super().__init__(root, prepared_dire=True)
+        
+    def __getitem__(self, idx):
+        img_path, dire_path, eps_path, isfake = self.img_paths[idx]
+    
+        dire = Image.open(dire_path).convert('RGB')
+        dire = TF.to_tensor(dire)*2 - 1
+        
+        return (dire, isfake), (dire_path,)
+        

datasets/README.md

@@ -0,0 +1 @@
+Please locate your dataset folder inside this directory.

deploy_utils/check_stat.py

@@ -0,0 +1,63 @@
+# create image resolution histogram
+
+from glob import glob 
+from PIL import Image
+from tqdm.auto import tqdm 
+
+
+tmfiles_f = glob("/truemedia-eval/images/fakes/*")
+tmfiles_r = glob("/truemedia-eval/images/reals/*")
+
+files_f = glob("/home/ubuntu/Datasets/diffusiondb/train/*")
+files_r = glob("/home/ubuntu/Datasets/ffhq/in-the-wild-images/train/*")
+files_f += glob("/home/ubuntu/Datasets/stylegan2-ffhq/train/*")
+files_r += glob("/home/ubuntu/Datasets/coco-2017-train/train2017/train/*")
+Ws_f = []
+Hs_f = []
+for f in tqdm(files_f):
+    img = Image.open(f)
+    w, h = img.size
+    Ws_f.append(w)
+    Hs_f.append(h)
+
+Ws_r = []
+Hs_r = []
+for f in tqdm(files_r):
+    img = Image.open(f)
+    w, h = img.size
+    Ws_r.append(w)
+    Hs_r.append(h)
+
+TMWs_r = []
+TMHs_r = []
+for f in tqdm(tmfiles_r):
+    img = Image.open(f)
+    w, h = img.size
+    TMWs_r.append(w)
+    TMHs_r.append(h)
+
+TMWs_f = []
+TMHs_f = []
+for f in tqdm(tmfiles_f):
+    img = Image.open(f)
+    w, h = img.size
+    TMWs_f.append(w)
+    TMHs_f.append(h)
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+# plot 2d (w, h)
+plt.figure(figsize=(8, 10))
+plt.scatter(Ws_r, Hs_r, s=3, alpha=0.5)
+plt.scatter(Ws_f, Hs_f, s=3, alpha=0.5, c='g')
+plt.scatter(TMWs_r, TMHs_r, s=3, alpha=0.5, c='r')
+plt.scatter(TMWs_f, TMHs_f, s=3, alpha=0.5, c='y')
+
+
+plt.xlabel('Width')
+plt.ylabel('Height')
+plt.legend(['Training-real', 'Training-fake', 'TrueMedia-real', 'TrueMedia-fake'])
+# save 
+plt.savefig('resolution_hist.png')
+

deploy_utils/check_stat_comp.py

@@ -0,0 +1,74 @@
+# create image resolution histogram
+
+from glob import glob 
+from PIL import Image
+from tqdm.auto import tqdm 
+import os
+
+tmfiles_f = glob("/home/ubuntu/y1/DistilDIRE/datasets/truemedia-political/images/fakes/*")
+tmfiles_r = glob("/home/ubuntu/y1/DistilDIRE/datasets/truemedia-political/images/reals/*")
+
+files_f = glob("/home/ubuntu/y1/DistilDIRE/datasets/y1-global-truemedia/images/fakes/*")
+files_r = glob("/truemedia-eval/y1dataset/images/reals/*")
+# files_f += glob("/home/ubuntu/Datasets/stylegan2-ffhq/train/*")
+# files_r += glob("/home/ubuntu/Datasets/coco-2017-train/train2017/train/*")
+training_f = []
+fake_cnt = 0
+real_cnt = 0
+for f in tqdm(files_f):
+    b = os.stat(f).st_size
+    img = Image.open(f)
+    w, h = img.size
+    if (b/(w*h)) < 1:
+        fake_cnt += 1
+    training_f.append(b/(w*h))
+
+training_r = []
+for f in tqdm(files_r):
+    try:
+        img = Image.open(f)
+        b = os.stat(f).st_size
+        w, h = img.size
+        if (b/(w*h)) < 1:
+            real_cnt += 1
+        training_r.append(b/(w*h))
+    except:
+        continue
+
+TM_r = []
+for f in tqdm(tmfiles_r):
+    try:
+        img = Image.open(f)
+        b = os.stat(f).st_size
+        w, h = img.size
+        TM_r.append(b/(w*h))
+    except:
+        continue
+
+TM_f = []
+for f in tqdm(tmfiles_f):
+    try:
+        img = Image.open(f)
+        b = os.stat(f).st_size
+        w, h = img.size
+        TM_f.append(b/(w*h))
+    except:
+        continue
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+# plot 2d (w, h)
+plt.figure(figsize=(8, 6))
+plt.hist(training_r, bins=100, alpha=0.5)
+plt.hist(training_f, bins=100, alpha=0.5, color='g')
+plt.hist(TM_r, bins=100, alpha=0.5, color='r')
+plt.hist(TM_f, bins=100, alpha=0.5, color='y')
+
+# Bytes/pixel 
+plt.xlabel('Bytes/pixel')
+plt.legend(['Training-real', 'Training-fake', 'TrueMedia-real', 'TrueMedia-fake'])
+# save 
+plt.savefig('comp_hist.png')
+print(f"Training real: {real_cnt}, Training fake: {fake_cnt}")
+

deploy_utils/concat.py

@@ -0,0 +1,32 @@
+import shutil
+from tqdm.auto import tqdm
+import os.path as osp
+from glob import glob
+from PIL import Image
+
+# PATH = "/home/ubuntu/Datasets/"
+ROOTS = ['/home/ubuntu/Datasets/coco-2017-train/train2017/train', "/home/ubuntu/Datasets/ms-5kimages", "/home/ubuntu/Datasets/ffhq/in-the-wild-images/train"]
+names = ['coco', 'ms-5k', 'ffhq']
+DEST = "/home/ubuntu/y1/y1-global-truemedia/train/images"
+skipped = 0
+for name, root in zip(names, ROOTS):
+    # cur_path = osp.join(root, 'train/images')
+    # fakes = glob(osp.join(cur_path, 'fakes', '*'))
+    reals = glob(osp.join(root, '*'))
+    # for fake in fakes:
+    #     fname = f'{root}_{osp.basename(fake)}'
+    #     if not osp.exists(osp.join(DEST, 'fakes', fname)):
+    #         shutil.copy(fake, osp.join(DEST, 'fakes', fname))
+    if name == 'coco':
+        reals = reals[:35000]
+    for real in tqdm(reals):
+        fname = f'{name}_{osp.basename(real)}'
+        size = Image.open(real).size
+        if max(size) > 2000:
+            skipped += 1
+            continue
+        if not osp.exists(osp.join(DEST, 'reals', fname)):
+            shutil.copy(real, osp.join(DEST, 'reals', fname))
+
+print(f"Skipped {skipped} images for being too large.")
+        

deploy_utils/crawl.py

@@ -0,0 +1,68 @@
+import os
+import requests
+from bs4 import BeautifulSoup
+from urllib.parse import urljoin
+from urllib.request import urlretrieve
+from tqdm.auto import tqdm
+from selenium import webdriver
+from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.common.by import By
+import time
+
+# Function to fetch and parse the webpage
+def fetch_images(url, cls_name):
+    options = Options()
+    options.add_argument("--headless")
+    options.add_argument("window-size=1400,1500")
+    options.add_argument("--disable-gpu")
+    options.add_argument("--no-sandbox")
+    options.add_argument("start-maximized")
+    options.add_argument("enable-automation")
+    options.add_argument("--disable-infobars")
+    options.add_argument("--disable-dev-shm-usage")
+    user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+    options.add_argument(f"user-agent={user_agent}")
+
+    driver = webdriver.Chrome(options=options)
+    driver.get(url)
+    time.sleep(5)
+    img_elements = driver.find_elements(By.CLASS_NAME, cls_name)
+    urls = []
+    for img in img_elements:
+        src_url = img.get_attribute('src')
+        urls.append(src_url)
+    return urls
+
+
+# Function to download images
+def download_images(img_urls, download_folder='downloaded_images', desc=""):
+    if not os.path.exists(download_folder):
+        os.makedirs(download_folder)
+
+    for img_url in tqdm(img_urls, desc=desc):
+        fname = img_url.split('/')[-1]
+        fname = fname.split('?')[0]
+        filename = os.path.join(download_folder, fname)
+        urlretrieve(img_url, filename)
+        print(f"Downloaded: {filename}")
+
+# Main function to coordinate the image crawling
+def main():
+    page_n=0
+    while True:
+        page_n += 1
+        url = f'https://www.freepik.com/search?ai=only&format=search&last_filter=page&last_value={page_n}&page={page_n}&query=people%2C+political'
+        url = f'https://app.all-images.ai/en/f/images/search?pageIndex={page_n}&pageSize=40&search=man'
+        cls_name = "_1286nb17"
+        download_path = "/truemedia-eval/crawled-fakes/images/fakes"
+        try:
+            img_urls = fetch_images(url, cls_name)
+            if len(img_urls) == 0:
+                break
+            download_images(img_urls, download_path, desc=f"Page {page_n}")
+            
+        except:
+            break
+
+if __name__ == "__main__":
+    main()

deploy_utils/crawl_civt.py

@@ -0,0 +1,103 @@
+import os
+import requests
+from bs4 import BeautifulSoup
+from urllib.parse import urljoin
+import urllib.request
+from tqdm.auto import tqdm
+from selenium import webdriver
+from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.common.by import By
+import time
+
+# Function to fetch and parse the webpage
+def fetch_images(url):
+    options = Options()
+    options.add_argument("--headless")
+    options.add_argument("window-size=1400,1500")
+    options.add_argument("--disable-gpu")
+    options.add_argument("--no-sandbox")
+    options.add_argument("start-maximized")
+    options.add_argument("enable-automation")
+    options.add_argument("--disable-infobars")
+    options.add_argument("--disable-dev-shm-usage")
+    user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+    options.add_argument(f"user-agent={user_agent}")
+
+    driver = webdriver.Chrome(options=options)
+    driver.set_window_size(1400, 1500)
+    driver.get(url)
+    time.sleep(5)
+    last_height = driver.execute_script("return document.scrollingElement.scrollHeight")
+    urls = []
+    while True :
+        # img_elements = driver.find_elements(By.CLASS_NAME, cls_name)
+        img_elements = driver.find_elements(By.TAG_NAME, 'img')
+        for img in img_elements:
+            urls.append(img.get_attribute('src') or img.get_attribute('data-src'))
+        driver.execute_script(f"document.scrollingElement.scrollTo(0, document.scrollingElement.scrollHeight);")
+        time.sleep(5)
+        new_height = driver.execute_script("return document.scrollingElement.scrollHeight")
+        print(f"last_height: {last_height}, new_height: {new_height}")
+        if new_height == last_height:
+            break
+            
+        last_height = new_height
+    
+    print(f"Total images: {len(urls)}")
+    
+    return urls
+
+
+# Function to download images
+def download_images(img_urls, download_folder='downloaded_images', desc=""):
+    if not os.path.exists(download_folder):
+        os.makedirs(download_folder)
+    headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'}
+
+    for img_url in tqdm(img_urls, desc=desc):
+        fname = img_url.split('/')[-1]
+        fname = fname.split('?')[0]
+        filename = os.path.join(download_folder, fname)
+        if os.path.exists(filename):
+            continue
+        req = urllib.request.Request(img_url, headers=headers)
+        with urllib.request.urlopen(req) as response, open(filename, 'wb') as out_file:
+            data = response.read()
+            out_file.write(data)
+
+
+
+def crawl(query):
+    url = f'https://openart.ai/search/{query}?method=similarity'
+    download_path = "/truemedia-eval/crawled-fakes/images/fakes"
+    try:
+        img_urls = fetch_images(url)
+        if len(img_urls) == 0:
+            return
+        download_images(img_urls, download_path, desc=f"query: {query}")
+    except:
+        return
+
+# Main function to coordinate the image crawling
+def main():
+    # queries = ['xi', 'trump', 'biden', 'putin', 'zelensky', 'gaza', 'israel', 'taylor', 'palesti', 'nsfw', 'instagram']
+    queries = ['xi', 'trump', 'biden', 'putin', 'zelensky', 'gaza', 'israel', 'taylor', 'palesti', 'nsfw', 'instagram',
+    "Hamas", "Hezbollah", "Iran",
+    "North Korea", "Kim Jong-un", "Uyghurs", "Hong Kong", "Taiwan", "Tibet", "Kashmir", "Syria", "Assad", "ISIS",
+    "Taliban", "Afghanistan", "Saudi", "Yemen", "Qatar", "Brexit", "Catalonia", "Black lives matter", "Antifa", "Proud Boys",
+    "Abortion", "Guns", "Second Amendment", "Police", "Execution", "Drugs", "Opioids", "Climate", "Paris",
+    "Green", "Keystone", "Fracking", "GMOs", "Net", "WikiLeaks", "Assange", "Snowden", "NSA", "Patriot",
+    "COVID", "Vaccines", "Masks", "5G", "QAnon", "Fraud", "Capitol", "Impeachment", "Hunter", "Epstein",
+    "Trafficking", "MeToo", "Weinstein", "Cancel", "Race", "Affirmative", "Transgender", "Marriage",
+    "Religion", "Nationalism", "Neo-Nazism", "Confederate", "Antisemitism", "Islamophobia", "Wall", "ICE",
+    "DACA", "Dreamers", "Separation", "Sanctuary", "Refugees", "Travel", "Benghazi", "Emails", "Russia",
+    "Mueller", "Ukraine", "Mar-a-Lago", "Taxes", "Emoluments", "Gerrymandering", "Voter", "Mail", "Citizens",
+    "PACs", "Breonna", "Floyd", "Rittenhouse", "Insurrection", "Ivanka", "Kushner", "Bannon", "Flynn", "Stone"
+]
+
+    for query in queries:
+        crawl(query)
+    
+
+if __name__ == "__main__":
+    main()

deploy_utils/down.py

@@ -0,0 +1,16 @@
+import pandas as pd 
+from tqdm.auto import tqdm 
+import subprocess 
+import os.path as osp
+
+df = pd.read_csv('/home/ubuntu/y1/truemedia.csv')
+for i in tqdm(range(len(df))):
+    line = df.iloc[i]
+    ID = line['Id']
+    fake = 'fake' in line['Ground Truth'].lower()
+    if osp.exists(f'/truemedia-eval/images/fakes/{ID}') or osp.exists(f'/truemedia-eval/images/reals/{ID}'):
+        continue
+    if fake:
+        subprocess.run(f"aws s3 cp s3://truemedia-media/{ID} /truemedia-eval/images/fakes/{ID}", shell=True)
+    else:
+        subprocess.run(f"aws s3 cp s3://truemedia-media/{ID} /truemedia-eval/images/reals/{ID}", shell=True)

deploy_utils/instaimages_format.py

@@ -0,0 +1,18 @@
+from shutil import copy 
+from glob import glob 
+from PIL import Image
+from tqdm.auto import tqdm
+import os 
+
+os.makedirs("/truemedia-eval/instagram/images/reals", exist_ok=True)
+files = glob("/truemedia-eval/instagram/**/*.jpg", recursive=True)
+for f in tqdm(files):
+    try:
+        img = Image.open(f)
+    except:
+        continue
+    size = img.size
+    if max(size) > 2000:
+        continue 
+    fname = f.split("/")[-1]
+    copy(f, f"/truemedia-eval/instagram/images/reals/{fname}")

guided_diffusion/LICENSE

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

guided_diffusion/README.md

@@ -0,0 +1,176 @@
+# guided-diffusion
+
+This is the codebase for [Diffusion Models Beat GANS on Image Synthesis](http://arxiv.org/abs/2105.05233).
+
+This repository is based on [openai/improved-diffusion](https://github.com/openai/improved-diffusion), with modifications for classifier conditioning and architecture improvements.
+
+# Download pre-trained models
+
+We have released checkpoints for the main models in the paper. Before using these models, please review the corresponding [model card](model-card.md) to understand the intended use and limitations of these models.
+
+Here are the download links for each model checkpoint:
+
+ * 64x64 classifier: [64x64_classifier.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/64x64_classifier.pt)
+ * 64x64 diffusion: [64x64_diffusion.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/64x64_diffusion.pt)
+ * 128x128 classifier: [128x128_classifier.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/128x128_classifier.pt)
+ * 128x128 diffusion: [128x128_diffusion.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/128x128_diffusion.pt)
+ * 256x256 classifier: [256x256_classifier.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_classifier.pt)
+ * 256x256 diffusion: [256x256_diffusion.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion.pt)
+ * 256x256 diffusion (not class conditional): [256x256_diffusion_uncond.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt)
+ * 512x512 classifier: [512x512_classifier.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/512x512_classifier.pt)
+ * 512x512 diffusion: [512x512_diffusion.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/512x512_diffusion.pt)
+ * 64x64 -> 256x256 upsampler: [64_256_upsampler.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/64_256_upsampler.pt)
+ * 128x128 -> 512x512 upsampler: [128_512_upsampler.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/128_512_upsampler.pt)
+ * LSUN bedroom: [lsun_bedroom.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/lsun_bedroom.pt)
+ * LSUN cat: [lsun_cat.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/lsun_cat.pt)
+ * LSUN horse: [lsun_horse.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/lsun_horse.pt)
+ * LSUN horse (no dropout): [lsun_horse_nodropout.pt](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/lsun_horse_nodropout.pt)
+
+# Sampling from pre-trained models
+
+To sample from these models, you can use the `classifier_sample.py`, `image_sample.py`, and `super_res_sample.py` scripts.
+Here, we provide flags for sampling from all of these models.
+We assume that you have downloaded the relevant model checkpoints into a folder called `models/`.
+
+For these examples, we will generate 100 samples with batch size 4. Feel free to change these values.
+
+```
+SAMPLE_FLAGS="--batch_size 4 --num_samples 100 --timestep_respacing 250"
+```
+
+## Classifier guidance
+
+Note for these sampling runs that you can set `--classifier_scale 0` to sample from the base diffusion model.
+You may also use the `image_sample.py` script instead of `classifier_sample.py` in that case.
+
+ * 64x64 model:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --diffusion_steps 1000 --dropout 0.1 --image_size 64 --learn_sigma True --noise_schedule cosine --num_channels 192 --num_head_channels 64 --num_res_blocks 3 --resblock_updown True --use_new_attention_order True --use_fp16 True --use_scale_shift_norm True"
+python classifier_sample.py $MODEL_FLAGS --classifier_scale 1.0 --classifier_path models/64x64_classifier.pt --classifier_depth 4 --model_path models/64x64_diffusion.pt $SAMPLE_FLAGS
+```
+
+ * 128x128 model:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --diffusion_steps 1000 --image_size 128 --learn_sigma True --noise_schedule linear --num_channels 256 --num_heads 4 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python classifier_sample.py $MODEL_FLAGS --classifier_scale 0.5 --classifier_path models/128x128_classifier.pt --model_path models/128x128_diffusion.pt $SAMPLE_FLAGS
+```
+
+ * 256x256 model:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python classifier_sample.py $MODEL_FLAGS --classifier_scale 1.0 --classifier_path models/256x256_classifier.pt --model_path models/256x256_diffusion.pt $SAMPLE_FLAGS
+```
+
+ * 256x256 model (unconditional):
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python classifier_sample.py $MODEL_FLAGS --classifier_scale 10.0 --classifier_path models/256x256_classifier.pt --model_path models/256x256_diffusion_uncond.pt $SAMPLE_FLAGS
+```
+
+ * 512x512 model:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --diffusion_steps 1000 --image_size 512 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 False --use_scale_shift_norm True"
+python classifier_sample.py $MODEL_FLAGS --classifier_scale 4.0 --classifier_path models/512x512_classifier.pt --model_path models/512x512_diffusion.pt $SAMPLE_FLAGS
+```
+
+## Upsampling
+
+For these runs, we assume you have some base samples in a file `64_samples.npz` or `128_samples.npz` for the two respective models.
+
+ * 64 -> 256:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --diffusion_steps 1000 --large_size 256  --small_size 64 --learn_sigma True --noise_schedule linear --num_channels 192 --num_heads 4 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python super_res_sample.py $MODEL_FLAGS --model_path models/64_256_upsampler.pt --base_samples 64_samples.npz $SAMPLE_FLAGS
+```
+
+ * 128 -> 512:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16 --class_cond True --diffusion_steps 1000 --large_size 512 --small_size 128 --learn_sigma True --noise_schedule linear --num_channels 192 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python super_res_sample.py $MODEL_FLAGS --model_path models/128_512_upsampler.pt $SAMPLE_FLAGS --base_samples 128_samples.npz
+```
+
+## LSUN models
+
+These models are class-unconditional and correspond to a single LSUN class. Here, we show how to sample from `lsun_bedroom.pt`, but the other two LSUN checkpoints should work as well:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --dropout 0.1 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python image_sample.py $MODEL_FLAGS --model_path models/lsun_bedroom.pt $SAMPLE_FLAGS
+```
+
+You can sample from `lsun_horse_nodropout.pt` by changing the dropout flag:
+
+```
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --dropout 0.0 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+python image_sample.py $MODEL_FLAGS --model_path models/lsun_horse_nodropout.pt $SAMPLE_FLAGS
+```
+
+Note that for these models, the best samples result from using 1000 timesteps:
+
+```
+SAMPLE_FLAGS="--batch_size 4 --num_samples 100 --timestep_respacing 1000"
+```
+
+# Results
+
+This table summarizes our ImageNet results for pure guided diffusion models:
+
+| Dataset          | FID  | Precision | Recall |
+|------------------|------|-----------|--------|
+| ImageNet 64x64   | 2.07 | 0.74      | 0.63   |
+| ImageNet 128x128 | 2.97 | 0.78      | 0.59   |
+| ImageNet 256x256 | 4.59 | 0.82      | 0.52   |
+| ImageNet 512x512 | 7.72 | 0.87      | 0.42   |
+
+This table shows the best results for high resolutions when using upsampling and guidance together:
+
+| Dataset          | FID  | Precision | Recall |
+|------------------|------|-----------|--------|
+| ImageNet 256x256 | 3.94 | 0.83      | 0.53   |
+| ImageNet 512x512 | 3.85 | 0.84      | 0.53   |
+
+Finally, here are the unguided results on individual LSUN classes:
+
+| Dataset      | FID  | Precision | Recall |
+|--------------|------|-----------|--------|
+| LSUN Bedroom | 1.90 | 0.66      | 0.51   |
+| LSUN Cat     | 5.57 | 0.63      | 0.52   |
+| LSUN Horse   | 2.57 | 0.71      | 0.55   |
+
+# Training models
+
+Training diffusion models is described in the [parent repository](https://github.com/openai/improved-diffusion). Training a classifier is similar. We assume you have put training hyperparameters into a `TRAIN_FLAGS` variable, and classifier hyperparameters into a `CLASSIFIER_FLAGS` variable. Then you can run:
+
+```
+mpiexec -n N python scripts/classifier_train.py --data_dir path/to/imagenet $TRAIN_FLAGS $CLASSIFIER_FLAGS
+```
+
+Make sure to divide the batch size in `TRAIN_FLAGS` by the number of MPI processes you are using.
+
+Here are flags for training the 128x128 classifier. You can modify these for training classifiers at other resolutions:
+
+```sh
+TRAIN_FLAGS="--iterations 300000 --anneal_lr True --batch_size 256 --lr 3e-4 --save_interval 10000 --weight_decay 0.05"
+CLASSIFIER_FLAGS="--image_size 128 --classifier_attention_resolutions 32,16,8 --classifier_depth 2 --classifier_width 128 --classifier_pool attention --classifier_resblock_updown True --classifier_use_scale_shift_norm True"
+```
+
+For sampling from a 128x128 classifier-guided model, 25 step DDIM:
+
+```sh
+MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond True --image_size 128 --learn_sigma True --num_channels 256 --num_heads 4 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
+CLASSIFIER_FLAGS="--image_size 128 --classifier_attention_resolutions 32,16,8 --classifier_depth 2 --classifier_width 128 --classifier_pool attention --classifier_resblock_updown True --classifier_use_scale_shift_norm True --classifier_scale 1.0 --classifier_use_fp16 True"
+SAMPLE_FLAGS="--batch_size 4 --num_samples 50000 --timestep_respacing ddim25 --use_ddim True"
+mpiexec -n N python scripts/classifier_sample.py \
+    --model_path /path/to/model.pt \
+    --classifier_path path/to/classifier.pt \
+    $MODEL_FLAGS $CLASSIFIER_FLAGS $SAMPLE_FLAGS
+```
+
+To sample for 250 timesteps without DDIM, replace `--timestep_respacing ddim25` to `--timestep_respacing 250`, and replace `--use_ddim True` with `--use_ddim False`.

guided_diffusion/compute_dire_eps.py

@@ -0,0 +1,266 @@
+"""
+Modified from guided-diffusion/scripts/image_sample.py
+"""
+
+import argparse
+import os
+import torch
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+from dataset import TMDistilDireDataset
+from torch.utils.data import DataLoader
+import cv2
+
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from tqdm.auto import tqdm 
+
+import numpy as np
+import torch as th
+import torchvision
+import os.path as osp
+
+
+from .guided_diffusion.script_util import (
+    NUM_CLASSES,
+    model_and_diffusion_defaults,
+    create_model_and_diffusion,
+    add_dict_to_argparser,
+    dict_parse,
+    args_to_dict,
+
+)
+
+
+
+def reshape_image(imgs: torch.Tensor, image_size: int) -> torch.Tensor:
+    if len(imgs.shape) == 3:
+        imgs = imgs.unsqueeze(0)
+    if imgs.shape[2] != imgs.shape[3]:
+        crop_func = transforms.CenterCrop(image_size)
+        imgs = crop_func(imgs)
+    if imgs.shape[2] != image_size:
+        imgs = F.interpolate(imgs, size=(image_size, image_size), mode="bicubic", antialias=True)
+    return imgs
+
+
+
+def create_argparser():
+    defaults = dict(
+        clip_denoised=True,
+        num_samples=-1,
+        use_ddim=True,
+        real_step=0,
+        continue_reverse=False,
+        has_subfolder=False,
+    )
+    defaults.update(model_and_diffusion_defaults())
+    sanic_dict = dict(
+       attention_resolutions='32,16,8',
+       class_cond=False,
+       diffusion_steps=1000,
+       image_size=256,
+       learn_sigma=True,
+       model_path="./models/256x256-adm.pt",
+       noise_schedule='linear',
+       num_channels=256,
+       num_head_channels=64,
+       num_res_blocks=2,
+       resblock_updown=True,
+       use_fp16=True,
+       use_scale_shift_norm=True,
+       data_root="",
+       compute_dire=False,
+       compute_eps=False,
+       save_root="",
+       batch_size=32,
+    )
+    defaults.update(sanic_dict)
+    parser = argparse.ArgumentParser()
+    add_dict_to_argparser(parser, defaults)
+    args = parser.parse_args()
+
+    return args_to_dict(args, list(defaults.keys()))
+
+def create_dicts_for_static_init():
+    defaults = dict(
+        clip_denoised=True,
+        num_samples=-1,
+        use_ddim=True,
+        real_step=0,
+        continue_reverse=False,
+        has_subfolder=False,
+    )
+    defaults.update(model_and_diffusion_defaults())
+    sanic_dict = dict(
+       attention_resolutions='32,16,8',
+       class_cond=False,
+       diffusion_steps=1000,
+       image_size=256,
+       learn_sigma=True,
+       model_path="./models/256x256-adm.pt",
+       noise_schedule='linear',
+       num_channels=256,
+       num_head_channels=64,
+       num_res_blocks=2,
+       resblock_updown=True,
+       use_fp16=True,
+       use_scale_shift_norm=True,
+       data_root="",
+       compute_dire=False,
+       compute_eps=False,
+       save_root="",
+       batch_size=32,
+    )
+    defaults.update(sanic_dict)
+
+    return defaults
+
+
+
+@torch.no_grad()
+def dire(img_batch:torch.Tensor, model, diffusion, args, save_img=False, save_path=None):
+    print("computing recons & DIRE ...")
+    imgs = img_batch.cuda()
+    
+    batch_size = imgs.shape[0]
+    model_kwargs = {}
+
+    reverse_fn = diffusion.ddim_reverse_sample_loop
+    assert (imgs.shape[2] == imgs.shape[3]) and (imgs.shape[3] == args['image_size']), f"Image size mismatch: {imgs.shape[2]} != {args['image_size']}"
+    latent = reverse_fn(
+        model,
+        (batch_size, 3, args['image_size'], args['image_size']),
+        noise=imgs,
+        clip_denoised=args['clip_denoised'],
+        model_kwargs=model_kwargs,
+        real_step=args['real_step'],
+    )
+    sample_fn = diffusion.p_sample_loop if not args['use_ddim'] else diffusion.ddim_sample_loop
+    recons = sample_fn(
+        model,
+        (batch_size, 3, args['image_size'], args['image_size']),
+        noise=latent,
+        clip_denoised=args['clip_denoised'],
+        model_kwargs=model_kwargs,
+        real_step=args['real_step']
+    )
+
+    dire = th.abs(imgs - recons)
+    dire = (dire*255./2).clamp(0, 255).to(th.uint8)
+    dire = dire.contiguous() / 255.
+    dire = (dire).clamp(0, 1).to(th.float32)
+    
+    # scale imgs and recons
+    imgs = (imgs+1)*0.5
+    recons = (recons+1)*0.5
+
+    if save_img:
+        # save images
+        for i in range(len(img_batch)):
+            dire_img = dire[i].detach().cpu().numpy().transpose(1, 2, 0)
+            dire_img = cv2.cvtColor(dire_img, cv2.COLOR_RGB2BGR)
+            
+            dire_path = os.path.join(save_path, f"dire_{i}.png")
+            cv2.imwrite(dire_path, dire_img*255)
+
+    return dire, imgs, recons
+
+@torch.no_grad()
+def dire_get_first_step_noise(img_batch:torch.Tensor, model, diffusion, args, device):
+    imgs = img_batch.to(device)    
+    batch_size = imgs.shape[0]
+    model_kwargs = {}
+
+    reverse_fn = diffusion.ddim_reverse_sample_only_eps
+    assert (imgs.shape[2] == imgs.shape[3]) and (imgs.shape[3] == args['image_size']), f"Image size mismatch: {imgs.shape[2]} != {args['image_size']}"
+    
+    eps = reverse_fn(
+        model,
+        # (batch_size, 3, args['image_size'], args['image_size']),
+        x=imgs,
+        t=torch.zeros(imgs.shape[0],).long().to(device),
+        clip_denoised=args['clip_denoised'],
+        model_kwargs=model_kwargs,
+        eta=0.0
+        #real_step=args['real_step'],
+    )
+
+    return eps
+
+
+if __name__ == "__main__":
+    from torch.utils.data.distributed import DistributedSampler
+    import torch.distributed as dist
+    import os 
+    
+    dist.init_process_group(backend='nccl', init_method='env://')
+    
+    local_rank = int(os.environ['LOCAL_RANK']) 
+    torch.cuda.set_device(local_rank)
+    
+    # Set device for this process
+    device = torch.device("cuda") 
+
+    adm_args = create_argparser()
+    adm_args['timestep_respacing'] = 'ddim20'
+    adm_model, diffusion = create_model_and_diffusion(**dict_parse(adm_args, model_and_diffusion_defaults().keys()))
+    print(f"checkpoint: {adm_args['model_path']}")
+    print(f"model channel: {adm_args['num_channels']}, {adm_model.model_channels}")
+    adm_model.load_state_dict(torch.load(adm_args['model_path'], map_location="cpu"))
+    adm_model.to(device)
+
+    adm_model.convert_to_fp16()
+    adm_model.eval()
+
+    dataset = TMDistilDireDataset(adm_args['data_root'], prepared_dire=False)
+    sampler = DistributedSampler(dataset, shuffle=False)
+    os.makedirs(osp.join(adm_args['save_root'], 'images', 'fakes'), exist_ok=True)
+    os.makedirs(osp.join(adm_args['save_root'], 'images', 'reals'), exist_ok=True)
+    os.makedirs(osp.join(adm_args['save_root'], 'dire', 'fakes'), exist_ok=True)
+    os.makedirs(osp.join(adm_args['save_root'], 'dire', 'reals'), exist_ok=True)
+    os.makedirs(osp.join(adm_args['save_root'], 'eps', 'fakes'), exist_ok=True)
+    os.makedirs(osp.join(adm_args['save_root'], 'eps', 'reals'), exist_ok=True)
+    print(f"Dataset length: {len(dataset)}")
+    dataloader = DataLoader(dataset, batch_size=adm_args['batch_size'], num_workers=4, drop_last=False, pin_memory=True, sampler=sampler)#
+
+    for (img_batch, dire_batch, eps_batch, isfake_batch), (img_pathes, dire_pathes, eps_pathes) in tqdm(dataloader):
+        haveall=True 
+        for i in range(len(img_batch)):
+            basename = osp.basename(img_pathes[i])
+            isfake = isfake_batch[i]
+            
+            img_path = osp.join(adm_args['save_root'], 'images', 'fakes', basename) if isfake else osp.join(adm_args['save_root'], 'images', 'reals', basename)
+            dire_path = img_path.replace('/images/', '/dire/')
+            eps_path = img_path.replace('/images/', '/eps/').split('.')[0] + '.pt'
+            if (not osp.exists(img_pathes[i])) or (not osp.exists(dire_path)) or (not osp.exists(eps_path)):
+                haveall=False
+                break
+        if haveall:
+            continue
+        with torch.no_grad():
+            eps = None
+            img = (img_batch.detach().cpu()+1)*0.5
+            if adm_args['compute_eps']:
+                eps = dire_get_first_step_noise(img_batch, adm_model, diffusion, adm_args, device)
+                eps = eps.detach().cpu()
+
+            if adm_args['compute_dire']:
+                dire_img, img, recons = dire(img_batch, adm_model, diffusion, adm_args)
+                dire_img = dire_img.detach().cpu()
+                img = img.detach().cpu()
+            
+            for i in range(len(img_batch)):
+                basename = osp.basename(img_pathes[i])
+                isfake = isfake_batch[i]
+                
+                img_path = osp.join(adm_args['save_root'], 'images', 'fakes', basename) if isfake else osp.join(adm_args['save_root'], 'images', 'reals', basename)
+                dire_path = img_path.replace('/images/', '/dire/')
+                eps_path = img_path.replace('/images/', '/eps/').split('.')[0] + '.pt'
+                
+                if not osp.exists(img_path):
+                    torchvision.utils.save_image(img[i], img_path)
+                if not osp.exists(dire_path) and adm_args['compute_dire']:
+                    torchvision.utils.save_image(dire_img[i], dire_path)
+                if not osp.exists(eps_path) and eps is not None:
+                    torch.save(eps[i], eps_path)

guided_diffusion/guided_diffusion/__init__.py

@@ -0,0 +1,3 @@
+"""
+Codebase for "Improved Denoising Diffusion Probabilistic Models".
+"""

guided_diffusion/guided_diffusion/dist_util.py

@@ -0,0 +1,99 @@
+"""
+Helpers for distributed training.
+"""
+
+import io
+import os
+import socket
+
+import blobfile as bf
+from mpi4py import MPI
+import torch as th
+import torch.distributed as dist
+from . import logger
+# Change this to reflect your cluster layout.
+# The GPU for a given rank is (rank % GPUS_PER_NODE).
+GPUS_PER_NODE = 8
+
+SETUP_RETRY_COUNT = 3
+
+
+def setup_dist(devices=""):
+    """
+    Setup a distributed process group.
+    """
+    if dist.is_initialized():
+        return
+    devices_list = [] 
+    for device in devices.split(","):
+        devices_list.append(int(device))
+    # print(devices_list)
+    GPUS_PER_NODE = len(devices_list)
+    # print("GPUS_PER_NODE: ",GPUS_PER_NODE)
+    os.environ["CUDA_VISIBLE_DEVICES"] = f"{devices_list[MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE]}"
+    # print(f"rank: {MPI.COMM_WORLD.Get_rank()}, {os.environ['CUDA_VISIBLE_DEVICES']}")
+    comm = MPI.COMM_WORLD
+    backend = "gloo" if not th.cuda.is_available() else "nccl"
+
+    if backend == "gloo":
+        hostname = "localhost"
+    else:
+        hostname = socket.gethostbyname(socket.getfqdn())
+    os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
+    os.environ["RANK"] = str(comm.rank)
+    os.environ["WORLD_SIZE"] = str(comm.size)
+
+    port = comm.bcast(_find_free_port(), root=0)
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group(backend=backend, init_method="env://")
+
+
+def dev():
+    """
+    Get the device to use for torch.distributed.
+    """
+    if th.cuda.is_available():
+        return th.device(f"cuda")
+    return th.device("cpu")
+
+
+def load_state_dict(path, **kwargs):
+    """
+    Load a PyTorch file without redundant fetches across MPI ranks.
+    """
+    chunk_size = 2 ** 30  # MPI has a relatively small size limit
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        with bf.BlobFile(path, "rb") as f:
+            data = f.read()
+        num_chunks = len(data) // chunk_size
+        if len(data) % chunk_size:
+            num_chunks += 1
+        MPI.COMM_WORLD.bcast(num_chunks)
+        for i in range(0, len(data), chunk_size):
+            MPI.COMM_WORLD.bcast(data[i : i + chunk_size])
+    else:
+        num_chunks = MPI.COMM_WORLD.bcast(None)
+        data = bytes()
+        for _ in range(num_chunks):
+            data += MPI.COMM_WORLD.bcast(None)
+
+    return th.load(io.BytesIO(data), **kwargs)
+
+
+def sync_params(params):
+    """
+    Synchronize a sequence of Tensors across ranks from rank 0.
+    """
+    for p in params:
+        with th.no_grad():
+            dist.broadcast(p, 0)
+
+
+def _find_free_port():
+    try:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+    finally:
+        s.close()

guided_diffusion/guided_diffusion/fp16_util.py

@@ -0,0 +1,237 @@
+"""
+Helpers to train with 16-bit precision.
+"""
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from . import logger
+
+INITIAL_LOG_LOSS_SCALE = 20.0
+
+
+def convert_module_to_f16(l):
+    """
+    Convert primitive modules to float16.
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.half()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.half()
+
+
+def convert_module_to_f32(l):
+    """
+    Convert primitive modules to float32, undoing convert_module_to_f16().
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.float()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.float()
+
+
+def make_master_params(param_groups_and_shapes):
+    """
+    Copy model parameters into a (differently-shaped) list of full-precision
+    parameters.
+    """
+    master_params = []
+    for param_group, shape in param_groups_and_shapes:
+        master_param = nn.Parameter(
+            _flatten_dense_tensors(
+                [param.detach().float() for (_, param) in param_group]
+            ).view(shape)
+        )
+        master_param.requires_grad = True
+        master_params.append(master_param)
+    return master_params
+
+
+def model_grads_to_master_grads(param_groups_and_shapes, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    for master_param, (param_group, shape) in zip(
+        master_params, param_groups_and_shapes
+    ):
+        master_param.grad = _flatten_dense_tensors(
+            [param_grad_or_zeros(param) for (_, param) in param_group]
+        ).view(shape)
+
+
+def master_params_to_model_params(param_groups_and_shapes, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    # Without copying to a list, if a generator is passed, this will
+    # silently not copy any parameters.
+    for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes):
+        for (_, param), unflat_master_param in zip(
+            param_group, unflatten_master_params(param_group, master_param.view(-1))
+        ):
+            param.detach().copy_(unflat_master_param)
+
+
+def unflatten_master_params(param_group, master_param):
+    return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group])
+
+
+def get_param_groups_and_shapes(named_model_params):
+    named_model_params = list(named_model_params)
+    scalar_vector_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim <= 1],
+        (-1),
+    )
+    matrix_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim > 1],
+        (1, -1),
+    )
+    return [scalar_vector_named_params, matrix_named_params]
+
+
+def master_params_to_state_dict(
+    model, param_groups_and_shapes, master_params, use_fp16
+):
+    if use_fp16:
+        state_dict = model.state_dict()
+        for master_param, (param_group, _) in zip(
+            master_params, param_groups_and_shapes
+        ):
+            for (name, _), unflat_master_param in zip(
+                param_group, unflatten_master_params(param_group, master_param.view(-1))
+            ):
+                assert name in state_dict
+                state_dict[name] = unflat_master_param
+    else:
+        state_dict = model.state_dict()
+        for i, (name, _value) in enumerate(model.named_parameters()):
+            assert name in state_dict
+            state_dict[name] = master_params[i]
+    return state_dict
+
+
+def state_dict_to_master_params(model, state_dict, use_fp16):
+    if use_fp16:
+        named_model_params = [
+            (name, state_dict[name]) for name, _ in model.named_parameters()
+        ]
+        param_groups_and_shapes = get_param_groups_and_shapes(named_model_params)
+        master_params = make_master_params(param_groups_and_shapes)
+    else:
+        master_params = [state_dict[name] for name, _ in model.named_parameters()]
+    return master_params
+
+
+def zero_master_grads(master_params):
+    for param in master_params:
+        param.grad = None
+
+
+def zero_grad(model_params):
+    for param in model_params:
+        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
+        if param.grad is not None:
+            param.grad.detach_()
+            param.grad.zero_()
+
+
+def param_grad_or_zeros(param):
+    if param.grad is not None:
+        return param.grad.data.detach()
+    else:
+        return th.zeros_like(param)
+
+
+class MixedPrecisionTrainer:
+    def __init__(
+        self,
+        *,
+        model,
+        use_fp16=False,
+        fp16_scale_growth=1e-3,
+        initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE,
+    ):
+        self.model = model
+        self.use_fp16 = use_fp16
+        self.fp16_scale_growth = fp16_scale_growth
+
+        self.model_params = list(self.model.parameters())
+        self.master_params = self.model_params
+        self.param_groups_and_shapes = None
+        self.lg_loss_scale = initial_lg_loss_scale
+
+        if self.use_fp16:
+            self.param_groups_and_shapes = get_param_groups_and_shapes(
+                self.model.named_parameters()
+            )
+            self.master_params = make_master_params(self.param_groups_and_shapes)
+            self.model.convert_to_fp16()
+
+    def zero_grad(self):
+        zero_grad(self.model_params)
+
+    def backward(self, loss: th.Tensor):
+        if self.use_fp16:
+            loss_scale = 2 ** self.lg_loss_scale
+            (loss * loss_scale).backward()
+        else:
+            loss.backward()
+
+    def optimize(self, opt: th.optim.Optimizer):
+        if self.use_fp16:
+            return self._optimize_fp16(opt)
+        else:
+            return self._optimize_normal(opt)
+
+    def _optimize_fp16(self, opt: th.optim.Optimizer):
+        logger.logkv_mean("lg_loss_scale", self.lg_loss_scale)
+        model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params)
+        grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale)
+        if check_overflow(grad_norm):
+            self.lg_loss_scale -= 1
+            logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}")
+            zero_master_grads(self.master_params)
+            return False
+
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+
+        for p in self.master_params:
+            p.grad.mul_(1.0 / (2 ** self.lg_loss_scale))
+        opt.step()
+        zero_master_grads(self.master_params)
+        master_params_to_model_params(self.param_groups_and_shapes, self.master_params)
+        self.lg_loss_scale += self.fp16_scale_growth
+        return True
+
+    def _optimize_normal(self, opt: th.optim.Optimizer):
+        grad_norm, param_norm = self._compute_norms()
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+        opt.step()
+        return True
+
+    def _compute_norms(self, grad_scale=1.0):
+        grad_norm = 0.0
+        param_norm = 0.0
+        for p in self.master_params:
+            with th.no_grad():
+                param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2
+                if p.grad is not None:
+                    grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2
+        return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm)
+
+    def master_params_to_state_dict(self, master_params):
+        return master_params_to_state_dict(
+            self.model, self.param_groups_and_shapes, master_params, self.use_fp16
+        )
+
+    def state_dict_to_master_params(self, state_dict):
+        return state_dict_to_master_params(self.model, state_dict, self.use_fp16)
+
+
+def check_overflow(value):
+    return (value == float("inf")) or (value == -float("inf")) or (value != value)

guided_diffusion/guided_diffusion/gaussian_diffusion.py

@@ -0,0 +1,1038 @@
+"""
+This code started out as a PyTorch port of Ho et al's diffusion models:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
+
+Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
+"""
+
+import enum
+import math
+
+import numpy as np
+import torch as th
+
+from .nn import mean_flat
+from .losses import normal_kl, discretized_gaussian_log_likelihood
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    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].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    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(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here, and then adapted over time to further experimentation.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    :param model_mean_type: a ModelMeanType determining what the model outputs.
+    :param model_var_type: a ModelVarType determining how variance is output.
+    :param loss_type: a LossType determining the loss function to use.
+    :param rescale_timesteps: if True, pass floating point timesteps into the
+                              model so that they are always scaled like in the
+                              original paper (0 to 1000).
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type,
+        rescale_timesteps=False,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.rescale_timesteps = rescale_timesteps
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # log calculation clipped because the posterior variance is 0 at the
+        # beginning of the diffusion chain.
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev)
+            * np.sqrt(alphas)
+            / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        )
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(
+            self.log_one_minus_alphas_cumprod, t, x_start.shape
+        )
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+            * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, self._scale_timesteps(t), **model_kwargs)
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == ModelVarType.LEARNED:
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(
+                    self.posterior_log_variance_clipped, t, x.shape
+                )
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
+            if self.model_mean_type == ModelMeanType.START_X:
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(
+                x_start=pred_xstart, x_t=x, t=t
+            )
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert (
+            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        )
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _scale_timesteps(self, t):
+        if self.rescale_timesteps:
+            return t.float() * (1000.0 / self.num_timesteps)
+        return t
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
+        new_mean = (
+            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        )
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
+            x, self._scale_timesteps(t), **model_kwargs
+        )
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(
+            x_start=out["pred_xstart"], x_t=x, t=t
+        )
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(
+                cond_fn, out, x, t, model_kwargs=model_kwargs
+            )
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_next)
+            + th.sqrt(1 - alpha_bar_next) * eps
+        )
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+    
+    def ddim_reverse_sample_only_eps(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+
+       
+        return eps
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        return_intermediate=False,
+        real_step=0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = []
+        if return_intermediate:
+            step = int(self.num_timesteps//100) if self.num_timesteps>100 else 1
+            i = 0
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+            real_step=real_step
+        ):
+            if return_intermediate:
+                i = i+1
+                if i % step==0:
+                    final.append(sample["sample"])
+            else:
+                final = sample["sample"]
+        return final
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        real_step=0
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1] if not real_step else list(range(self.num_timesteps))[:real_step][::-1]
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_reverse_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        real_step=0
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_reverse_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+            real_step=real_step
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_reverse_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        real_step=0
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        # print(real_step, type(real_step))
+        indices = list(range(self.num_timesteps)) if real_step==0 else list(range(self.num_timesteps))[:real_step]#[::-1] # 1, ..., T
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_reverse_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+    ):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
+                    x_start=x_start, x_t=x_t, t=t
+                )[0],
+                ModelMeanType.START_X: x_start,
+                ModelMeanType.EPSILON: noise,
+            }[self.model_mean_type]
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)

guided_diffusion/guided_diffusion/image_datasets.py

@@ -0,0 +1,266 @@
+import math
+import os
+import random
+
+from PIL import Image
+import blobfile as bf
+from mpi4py import MPI
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from guided_diffusion import logger
+
+
+def load_data_for_reverse(
+    *,
+    data_dir,
+    batch_size,
+    image_size,
+    class_cond=False,
+    deterministic=True,
+    random_crop=False,
+    random_flip=False,
+):
+    """
+    For a dataset, create a generator over (images, kwargs) pairs.
+
+    Each images is an NCHW float tensor, and the kwargs dict contains zero or
+    more keys, each of which map to a batched Tensor of their own.
+    The kwargs dict can be used for class labels, in which case the key is "y"
+    and the values are integer tensors of class labels. # kwargs包含了标签信息
+
+    :param data_dir: a dataset directory.
+    :param batch_size: the batch size of each returned pair.
+    :param image_size: the size to which images are resized.
+    :param class_cond: if True, include a "y" key in returned dicts for class
+                       label. If classes are not available and this is true, an
+                       exception will be raised.
+    :param deterministic: if True, yield results in a deterministic order.
+    :param random_crop: if True, randomly crop the images for augmentation.
+    :param random_flip: if True, randomly flip the images for augmentation.
+    """
+    if not data_dir:
+        raise ValueError("unspecified data directory")
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        all_files = _list_image_files_recursively(data_dir)
+
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        MPI.COMM_WORLD.bcast(all_files)
+    else:
+        all_files = MPI.COMM_WORLD.bcast(None)
+
+    classes = None
+    if class_cond:
+        # Assume classes are the first part of the filename,
+        # before an underscore.
+        class_names = [bf.basename(path).split("_")[0] for path in all_files]  # 标签_文件名
+        sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}
+        classes = [sorted_classes[x] for x in class_names]
+    dataset = ImageDataset_for_reverse(
+        image_size,
+        all_files,
+        classes=classes,
+        shard=MPI.COMM_WORLD.Get_rank(),
+        num_shards=MPI.COMM_WORLD.Get_size(),
+        random_crop=random_crop,
+        random_flip=random_flip,
+    )
+    logger.log("dataset length: {}".format(dataset.__len__() * MPI.COMM_WORLD.size))
+    if deterministic:
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=False)
+    else:
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True)
+    while True:
+        yield from loader
+
+
+def load_data(
+    *,
+    data_dir,
+    batch_size,
+    image_size,
+    class_cond=False,
+    deterministic=False,
+    random_crop=False,
+    random_flip=True,
+):
+    """
+    For a dataset, create a generator over (images, kwargs) pairs.
+
+    Each images is an NCHW float tensor, and the kwargs dict contains zero or
+    more keys, each of which map to a batched Tensor of their own.
+    The kwargs dict can be used for class labels, in which case the key is "y"
+    and the values are integer tensors of class labels.
+
+    :param data_dir: a dataset directory.
+    :param batch_size: the batch size of each returned pair.
+    :param image_size: the size to which images are resized.
+    :param class_cond: if True, include a "y" key in returned dicts for class
+                       label. If classes are not available and this is true, an
+                       exception will be raised.
+    :param deterministic: if True, yield results in a deterministic order.
+    :param random_crop: if True, randomly crop the images for augmentation.
+    :param random_flip: if True, randomly flip the images for augmentation.
+    """
+    if not data_dir:
+        raise ValueError("unspecified data directory")
+    all_files = _list_image_files_recursively(data_dir)
+    classes = None
+    if class_cond:
+        # Assume classes are the first part of the filename,
+        # before an underscore.
+        class_names = [bf.basename(path).split("_")[0] for path in all_files]
+        sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}
+        classes = [sorted_classes[x] for x in class_names]
+    dataset = ImageDataset(
+        image_size,
+        all_files,
+        classes=classes,
+        shard=MPI.COMM_WORLD.Get_rank(),
+        num_shards=MPI.COMM_WORLD.Get_size(),
+        random_crop=random_crop,
+        random_flip=random_flip,
+    )
+    if deterministic:
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True)
+    else:
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True)
+    while True:
+        yield from loader
+
+
+def _list_image_files_recursively(data_dir):
+    results = []
+    for entry in sorted(bf.listdir(data_dir)):
+        full_path = bf.join(data_dir, entry)
+        ext = entry.split(".")[-1]
+        if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]:
+            results.append(full_path)
+        elif bf.isdir(full_path):
+            results.extend(_list_image_files_recursively(full_path))
+    return results
+
+
+class ImageDataset_for_reverse(Dataset):
+    def __init__(
+        self,
+        resolution,
+        image_paths,
+        classes=None,
+        shard=0,
+        num_shards=1,
+        random_crop=False,
+        random_flip=True,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.local_images = image_paths[shard:][::num_shards]
+        self.local_classes = None if classes is None else classes[shard:][::num_shards]
+        self.random_crop = random_crop
+        self.random_flip = random_flip
+
+    def __len__(self):
+        return len(self.local_images)
+
+    def __getitem__(self, idx):
+        path = self.local_images[idx]
+        with bf.BlobFile(path, "rb") as f:
+            pil_image = Image.open(f)
+            pil_image.load()
+        pil_image = pil_image.convert("RGB")
+
+        if self.random_crop:
+            arr = random_crop_arr(pil_image, self.resolution)
+        else:
+            arr = center_crop_arr(pil_image, self.resolution)
+
+        if self.random_flip and random.random() < 0.5:
+            arr = arr[:, ::-1]
+
+        arr = arr.astype(np.float32) / 127.5 - 1
+
+        out_dict = {}
+        if self.local_classes is not None:
+            out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
+        arr = torch.from_numpy(np.transpose(arr, [2, 0, 1]))
+        return arr, out_dict, path
+
+
+class ImageDataset(Dataset):
+    def __init__(
+        self,
+        resolution,
+        image_paths,
+        classes=None,
+        shard=0,
+        num_shards=1,
+        random_crop=False,
+        random_flip=True,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.local_images = image_paths[shard:][::num_shards]
+        self.local_classes = None if classes is None else classes[shard:][::num_shards]
+        self.random_crop = random_crop
+        self.random_flip = random_flip
+
+    def __len__(self):
+        return len(self.local_images)
+
+    def __getitem__(self, idx):
+        path = self.local_images[idx]
+        with bf.BlobFile(path, "rb") as f:
+            pil_image = Image.open(f)
+            pil_image.load()
+        pil_image = pil_image.convert("RGB")
+
+        if self.random_crop:
+            arr = random_crop_arr(pil_image, self.resolution)
+        else:
+            arr = center_crop_arr(pil_image, self.resolution)
+
+        if self.random_flip and random.random() < 0.5:
+            arr = arr[:, ::-1]
+
+        arr = arr.astype(np.float32) / 127.5 - 1
+
+        out_dict = {}
+        if self.local_classes is not None:
+            out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
+        return np.transpose(arr, [2, 0, 1]), out_dict
+
+
+def center_crop_arr(pil_image, image_size):
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
+
+
+def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
+    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
+    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
+    smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)
+
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * smaller_dim_size:
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+
+    scale = smaller_dim_size / min(*pil_image.size)
+    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+
+    arr = np.array(pil_image)
+    crop_y = random.randrange(arr.shape[0] - image_size + 1)
+    crop_x = random.randrange(arr.shape[1] - image_size + 1)
+    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]

guided_diffusion/guided_diffusion/logger.py

@@ -0,0 +1,495 @@
+"""
+Logger copied from OpenAI baselines to avoid extra RL-based dependencies:
+https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py
+"""
+
+import os
+import sys
+import shutil
+import os.path as osp
+import json
+import time
+import datetime
+import tempfile
+import warnings
+from collections import defaultdict
+from contextlib import contextmanager
+
+DEBUG = 10
+INFO = 20
+WARN = 30
+ERROR = 40
+
+DISABLED = 50
+
+
+class KVWriter(object):
+    def writekvs(self, kvs):
+        raise NotImplementedError
+
+
+class SeqWriter(object):
+    def writeseq(self, seq):
+        raise NotImplementedError
+
+
+class HumanOutputFormat(KVWriter, SeqWriter):
+    def __init__(self, filename_or_file):
+        if isinstance(filename_or_file, str):
+            self.file = open(filename_or_file, "wt")
+            self.own_file = True
+        else:
+            assert hasattr(filename_or_file, "read"), (
+                "expected file or str, got %s" % filename_or_file
+            )
+            self.file = filename_or_file
+            self.own_file = False
+
+    def writekvs(self, kvs):
+        # Create strings for printing
+        key2str = {}
+        for (key, val) in sorted(kvs.items()):
+            if hasattr(val, "__float__"):
+                valstr = "%-8.3g" % val
+            else:
+                valstr = str(val)
+            key2str[self._truncate(key)] = self._truncate(valstr)
+
+        # Find max widths
+        if len(key2str) == 0:
+            print("WARNING: tried to write empty key-value dict")
+            return
+        else:
+            keywidth = max(map(len, key2str.keys()))
+            valwidth = max(map(len, key2str.values()))
+
+        # Write out the data
+        dashes = "-" * (keywidth + valwidth + 7)
+        lines = [dashes]
+        for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()):
+            lines.append(
+                "| %s%s | %s%s |"
+                % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val)))
+            )
+        lines.append(dashes)
+        self.file.write("\n".join(lines) + "\n")
+
+        # Flush the output to the file
+        self.file.flush()
+
+    def _truncate(self, s):
+        maxlen = 30
+        return s[: maxlen - 3] + "..." if len(s) > maxlen else s
+
+    def writeseq(self, seq):
+        seq = list(seq)
+        for (i, elem) in enumerate(seq):
+            self.file.write(elem)
+            if i < len(seq) - 1:  # add space unless this is the last one
+                self.file.write(" ")
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        if self.own_file:
+            self.file.close()
+
+
+class JSONOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "wt")
+
+    def writekvs(self, kvs):
+        for k, v in sorted(kvs.items()):
+            if hasattr(v, "dtype"):
+                kvs[k] = float(v)
+        self.file.write(json.dumps(kvs) + "\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class CSVOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "w+t")
+        self.keys = []
+        self.sep = ","
+
+    def writekvs(self, kvs):
+        # Add our current row to the history
+        extra_keys = list(kvs.keys() - self.keys)
+        extra_keys.sort()
+        if extra_keys:
+            self.keys.extend(extra_keys)
+            self.file.seek(0)
+            lines = self.file.readlines()
+            self.file.seek(0)
+            for (i, k) in enumerate(self.keys):
+                if i > 0:
+                    self.file.write(",")
+                self.file.write(k)
+            self.file.write("\n")
+            for line in lines[1:]:
+                self.file.write(line[:-1])
+                self.file.write(self.sep * len(extra_keys))
+                self.file.write("\n")
+        for (i, k) in enumerate(self.keys):
+            if i > 0:
+                self.file.write(",")
+            v = kvs.get(k)
+            if v is not None:
+                self.file.write(str(v))
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class TensorBoardOutputFormat(KVWriter):
+    """
+    Dumps key/value pairs into TensorBoard's numeric format.
+    """
+
+    def __init__(self, dir):
+        os.makedirs(dir, exist_ok=True)
+        self.dir = dir
+        self.step = 1
+        prefix = "events"
+        path = osp.join(osp.abspath(dir), prefix)
+        import tensorflow as tf
+        from tensorflow.python import pywrap_tensorflow
+        from tensorflow.core.util import event_pb2
+        from tensorflow.python.util import compat
+
+        self.tf = tf
+        self.event_pb2 = event_pb2
+        self.pywrap_tensorflow = pywrap_tensorflow
+        self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path))
+
+    def writekvs(self, kvs):
+        def summary_val(k, v):
+            kwargs = {"tag": k, "simple_value": float(v)}
+            return self.tf.Summary.Value(**kwargs)
+
+        summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()])
+        event = self.event_pb2.Event(wall_time=time.time(), summary=summary)
+        event.step = (
+            self.step
+        )  # is there any reason why you'd want to specify the step?
+        self.writer.WriteEvent(event)
+        self.writer.Flush()
+        self.step += 1
+
+    def close(self):
+        if self.writer:
+            self.writer.Close()
+            self.writer = None
+
+
+def make_output_format(format, ev_dir, log_suffix=""):
+    os.makedirs(ev_dir, exist_ok=True)
+    if format == "stdout":
+        return HumanOutputFormat(sys.stdout)
+    elif format == "log":
+        return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix))
+    elif format == "json":
+        return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix))
+    elif format == "csv":
+        return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix))
+    elif format == "tensorboard":
+        return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix))
+    else:
+        raise ValueError("Unknown format specified: %s" % (format,))
+
+
+# ================================================================
+# API
+# ================================================================
+
+
+def logkv(key, val):
+    """
+    Log a value of some diagnostic
+    Call this once for each diagnostic quantity, each iteration
+    If called many times, last value will be used.
+    """
+    get_current().logkv(key, val)
+
+
+def logkv_mean(key, val):
+    """
+    The same as logkv(), but if called many times, values averaged.
+    """
+    get_current().logkv_mean(key, val)
+
+
+def logkvs(d):
+    """
+    Log a dictionary of key-value pairs
+    """
+    for (k, v) in d.items():
+        logkv(k, v)
+
+
+def dumpkvs():
+    """
+    Write all of the diagnostics from the current iteration
+    """
+    return get_current().dumpkvs()
+
+
+def getkvs():
+    return get_current().name2val
+
+
+def log(*args, level=INFO):
+    """
+    Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).
+    """
+    get_current().log(*args, level=level)
+
+
+def debug(*args):
+    log(*args, level=DEBUG)
+
+
+def info(*args):
+    log(*args, level=INFO)
+
+
+def warn(*args):
+    log(*args, level=WARN)
+
+
+def error(*args):
+    log(*args, level=ERROR)
+
+
+def set_level(level):
+    """
+    Set logging threshold on current logger.
+    """
+    get_current().set_level(level)
+
+
+def set_comm(comm):
+    get_current().set_comm(comm)
+
+
+def get_dir():
+    """
+    Get directory that log files are being written to.
+    will be None if there is no output directory (i.e., if you didn't call start)
+    """
+    return get_current().get_dir()
+
+
+record_tabular = logkv
+dump_tabular = dumpkvs
+
+
+@contextmanager
+def profile_kv(scopename):
+    logkey = "wait_" + scopename
+    tstart = time.time()
+    try:
+        yield
+    finally:
+        get_current().name2val[logkey] += time.time() - tstart
+
+
+def profile(n):
+    """
+    Usage:
+    @profile("my_func")
+    def my_func(): code
+    """
+
+    def decorator_with_name(func):
+        def func_wrapper(*args, **kwargs):
+            with profile_kv(n):
+                return func(*args, **kwargs)
+
+        return func_wrapper
+
+    return decorator_with_name
+
+
+# ================================================================
+# Backend
+# ================================================================
+
+
+def get_current():
+    if Logger.CURRENT is None:
+        _configure_default_logger()
+
+    return Logger.CURRENT
+
+
+class Logger(object):
+    DEFAULT = None  # A logger with no output files. (See right below class definition)
+    # So that you can still log to the terminal without setting up any output files
+    CURRENT = None  # Current logger being used by the free functions above
+
+    def __init__(self, dir, output_formats, comm=None):
+        self.name2val = defaultdict(float)  # values this iteration
+        self.name2cnt = defaultdict(int)
+        self.level = INFO
+        self.dir = dir
+        self.output_formats = output_formats
+        self.comm = comm
+
+    # Logging API, forwarded
+    # ----------------------------------------
+    def logkv(self, key, val):
+        self.name2val[key] = val
+
+    def logkv_mean(self, key, val):
+        oldval, cnt = self.name2val[key], self.name2cnt[key]
+        self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)
+        self.name2cnt[key] = cnt + 1
+
+    def dumpkvs(self):
+        if self.comm is None:
+            d = self.name2val
+        else:
+            d = mpi_weighted_mean(
+                self.comm,
+                {
+                    name: (val, self.name2cnt.get(name, 1))
+                    for (name, val) in self.name2val.items()
+                },
+            )
+            if self.comm.rank != 0:
+                d["dummy"] = 1  # so we don't get a warning about empty dict
+        out = d.copy()  # Return the dict for unit testing purposes
+        for fmt in self.output_formats:
+            if isinstance(fmt, KVWriter):
+                fmt.writekvs(d)
+        self.name2val.clear()
+        self.name2cnt.clear()
+        return out
+
+    def log(self, *args, level=INFO):
+        if self.level <= level:
+            self._do_log(args)
+
+    # Configuration
+    # ----------------------------------------
+    def set_level(self, level):
+        self.level = level
+
+    def set_comm(self, comm):
+        self.comm = comm
+
+    def get_dir(self):
+        return self.dir
+
+    def close(self):
+        for fmt in self.output_formats:
+            fmt.close()
+
+    # Misc
+    # ----------------------------------------
+    def _do_log(self, args):
+        for fmt in self.output_formats:
+            if isinstance(fmt, SeqWriter):
+                fmt.writeseq(map(str, args))
+
+
+def get_rank_without_mpi_import():
+    # check environment variables here instead of importing mpi4py
+    # to avoid calling MPI_Init() when this module is imported
+    for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]:
+        if varname in os.environ:
+            return int(os.environ[varname])
+    return 0
+
+
+def mpi_weighted_mean(comm, local_name2valcount):
+    """
+    Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110
+    Perform a weighted average over dicts that are each on a different node
+    Input: local_name2valcount: dict mapping key -> (value, count)
+    Returns: key -> mean
+    """
+    all_name2valcount = comm.gather(local_name2valcount)
+    if comm.rank == 0:
+        name2sum = defaultdict(float)
+        name2count = defaultdict(float)
+        for n2vc in all_name2valcount:
+            for (name, (val, count)) in n2vc.items():
+                try:
+                    val = float(val)
+                except ValueError:
+                    if comm.rank == 0:
+                        warnings.warn(
+                            "WARNING: tried to compute mean on non-float {}={}".format(
+                                name, val
+                            )
+                        )
+                else:
+                    name2sum[name] += val * count
+                    name2count[name] += count
+        return {name: name2sum[name] / name2count[name] for name in name2sum}
+    else:
+        return {}
+
+
+def configure(dir=None, format_strs=None, comm=None, log_suffix=""):
+    """
+    If comm is provided, average all numerical stats across that comm
+    """
+    if dir is None:
+        dir = os.getenv("OPENAI_LOGDIR")
+    if dir is None:
+        dir = osp.join(
+            tempfile.gettempdir(),
+            datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"),
+        )
+    assert isinstance(dir, str)
+    dir = os.path.expanduser(dir)
+    os.makedirs(os.path.expanduser(dir), exist_ok=True)
+
+    rank = get_rank_without_mpi_import()
+    if rank > 0:
+        log_suffix = log_suffix + "-rank%03i" % rank
+
+    if format_strs is None:
+        if rank == 0:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",")
+        else:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",")
+    format_strs = filter(None, format_strs)
+    output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs]
+
+    Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm)
+    if output_formats:
+        log("Logging to %s" % dir)
+
+
+def _configure_default_logger():
+    configure()
+    Logger.DEFAULT = Logger.CURRENT
+
+
+def reset():
+    if Logger.CURRENT is not Logger.DEFAULT:
+        Logger.CURRENT.close()
+        Logger.CURRENT = Logger.DEFAULT
+        log("Reset logger")
+
+
+@contextmanager
+def scoped_configure(dir=None, format_strs=None, comm=None):
+    prevlogger = Logger.CURRENT
+    configure(dir=dir, format_strs=format_strs, comm=comm)
+    try:
+        yield
+    finally:
+        Logger.CURRENT.close()
+        Logger.CURRENT = prevlogger
+

guided_diffusion/guided_diffusion/losses.py

@@ -0,0 +1,77 @@
+"""
+Helpers for various likelihood-based losses. These are ported from the original
+Ho et al. diffusion models codebase:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py
+"""
+
+import numpy as np
+
+import torch as th
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs

guided_diffusion/guided_diffusion/nn.py

@@ -0,0 +1,170 @@
+"""
+Various utilities for neural networks.
+"""
+
+import math
+
+import torch as th
+import torch.nn as nn
+
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+    def forward(self, x):
+        return x * th.sigmoid(x)
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def update_ema(target_params, source_params, rate=0.99):
+    """
+    Update target parameters to be closer to those of source parameters using
+    an exponential moving average.
+
+    :param target_params: the target parameter sequence.
+    :param source_params: the source parameter sequence.
+    :param rate: the EMA rate (closer to 1 means slower).
+    """
+    for targ, src in zip(target_params, source_params):
+        targ.detach().mul_(rate).add_(src, alpha=1 - rate)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNorm32(32, channels)
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = th.exp(
+        -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+
+class CheckpointFunction(th.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        with th.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with th.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = th.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads

guided_diffusion/guided_diffusion/resample.py

@@ -0,0 +1,154 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch as th
+import torch.distributed as dist
+
+
+def create_named_schedule_sampler(name, diffusion):
+    """
+    Create a ScheduleSampler from a library of pre-defined samplers.
+
+    :param name: the name of the sampler.
+    :param diffusion: the diffusion object to sample for.
+    """
+    if name == "uniform":
+        return UniformSampler(diffusion)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(diffusion)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {name}")
+
+
+class ScheduleSampler(ABC):
+    """
+    A distribution over timesteps in the diffusion process, intended to reduce
+    variance of the objective.
+
+    By default, samplers perform unbiased importance sampling, in which the
+    objective's mean is unchanged.
+    However, subclasses may override sample() to change how the resampled
+    terms are reweighted, allowing for actual changes in the objective.
+    """
+
+    @abstractmethod
+    def weights(self):
+        """
+        Get a numpy array of weights, one per diffusion step.
+
+        The weights needn't be normalized, but must be positive.
+        """
+
+    def sample(self, batch_size, device):
+        """
+        Importance-sample timesteps for a batch.
+
+        :param batch_size: the number of timesteps.
+        :param device: the torch device to save to.
+        :return: a tuple (timesteps, weights):
+                 - timesteps: a tensor of timestep indices.
+                 - weights: a tensor of weights to scale the resulting losses.
+        """
+        w = self.weights()
+        p = w / np.sum(w)
+        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
+        indices = th.from_numpy(indices_np).long().to(device)
+        weights_np = 1 / (len(p) * p[indices_np])
+        weights = th.from_numpy(weights_np).float().to(device)
+        return indices, weights
+
+
+class UniformSampler(ScheduleSampler):
+    def __init__(self, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()

guided_diffusion/guided_diffusion/respace.py

@@ -0,0 +1,128 @@
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(
+            model, self.timestep_map, self.rescale_timesteps, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        if self.rescale_timesteps:
+            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)

guided_diffusion/guided_diffusion/script_util.py

@@ -0,0 +1,456 @@
+import argparse
+import inspect
+
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+from .unet import SuperResModel, UNetModel, EncoderUNetModel
+
+NUM_CLASSES = 1000
+
+
+def diffusion_defaults():
+    """
+    Defaults for image and classifier training.
+    """
+    return dict(
+        learn_sigma=False,
+        diffusion_steps=1000,
+        noise_schedule="linear",
+        timestep_respacing="",
+        use_kl=False,
+        predict_xstart=False,
+        rescale_timesteps=False,
+        rescale_learned_sigmas=False,
+    )
+
+
+
+def classifier_defaults():
+    """
+    Defaults for classifier models.
+    """
+    return dict(
+        image_size=64,
+        classifier_use_fp16=False,
+        classifier_width=128,
+        classifier_depth=2,
+        classifier_attention_resolutions="32,16,8",  # 16
+        classifier_use_scale_shift_norm=True,  # False
+        classifier_resblock_updown=True,  # False
+        classifier_pool="attention",
+    )
+
+def model_and_diffusion_defaults():
+    """
+    Defaults for image training.
+    """
+    res = dict(
+        image_size=64,
+        num_channels=128,
+        num_res_blocks=2,
+        num_heads=4,
+        num_heads_upsample=-1,
+        num_head_channels=-1,
+        attention_resolutions="16,8",
+        channel_mult="",
+        dropout=0.0,
+        class_cond=False,
+        use_checkpoint=False,
+        use_scale_shift_norm=True,
+        resblock_updown=False,
+        use_fp16=False,
+        use_new_attention_order=False,
+    )
+    res.update(diffusion_defaults())
+    return res
+
+
+
+def classifier_and_diffusion_defaults():
+    res = classifier_defaults()
+    res.update(diffusion_defaults())
+    return res
+
+
+def create_model_and_diffusion(
+    image_size,
+    class_cond,
+    learn_sigma,
+    num_channels,
+    num_res_blocks,
+    channel_mult,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    attention_resolutions,
+    dropout,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    use_checkpoint,
+    use_scale_shift_norm,
+    resblock_updown,
+    use_fp16,
+    use_new_attention_order,
+):
+    model = create_model(
+        image_size,
+        num_channels,
+        num_res_blocks,
+        channel_mult=channel_mult,
+        learn_sigma=learn_sigma,
+        class_cond=class_cond,
+        use_checkpoint=use_checkpoint,
+        attention_resolutions=attention_resolutions,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        dropout=dropout,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+        use_new_attention_order=use_new_attention_order,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return model, diffusion
+
+
+def create_model(
+    image_size,
+    num_channels,
+    num_res_blocks,
+    channel_mult="",
+    learn_sigma=False,
+    class_cond=False,
+    use_checkpoint=False,
+    attention_resolutions="16",
+    num_heads=1,
+    num_head_channels=-1,
+    num_heads_upsample=-1,
+    use_scale_shift_norm=False,
+    dropout=0,
+    resblock_updown=False,
+    use_fp16=False,
+    use_new_attention_order=False,
+):
+    if channel_mult == "":
+        if image_size == 512:
+            channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
+        elif image_size == 256:
+            channel_mult = (1, 1, 2, 2, 4, 4)
+        elif image_size == 128:
+            channel_mult = (1, 1, 2, 3, 4)
+        elif image_size == 64:
+            channel_mult = (1, 2, 3, 4)
+        else:
+            raise ValueError(f"unsupported image size: {image_size}")
+    else:
+        channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(","))
+
+    attention_ds = []
+    for res in attention_resolutions.split(","):
+        attention_ds.append(image_size // int(res))
+
+    return UNetModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=dropout,
+        channel_mult=channel_mult,
+        num_classes=(NUM_CLASSES if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        use_fp16=use_fp16,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_new_attention_order=use_new_attention_order,
+    )
+
+
+def create_classifier_and_diffusion(
+    image_size,
+    classifier_use_fp16,
+    classifier_width,
+    classifier_depth,
+    classifier_attention_resolutions,
+    classifier_use_scale_shift_norm,
+    classifier_resblock_updown,
+    classifier_pool,
+    learn_sigma,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+):
+    classifier = create_classifier(
+        image_size,
+        classifier_use_fp16,
+        classifier_width,
+        classifier_depth,
+        classifier_attention_resolutions,
+        classifier_use_scale_shift_norm,
+        classifier_resblock_updown,
+        classifier_pool,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return classifier, diffusion
+
+
+def create_classifier(
+    image_size,
+    classifier_use_fp16,
+    classifier_width,
+    classifier_depth,
+    classifier_attention_resolutions,
+    classifier_use_scale_shift_norm,
+    classifier_resblock_updown,
+    classifier_pool,
+):
+    if image_size == 512:
+        channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
+    elif image_size == 256:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif image_size == 128:
+        channel_mult = (1, 1, 2, 3, 4)
+    elif image_size == 64:
+        channel_mult = (1, 2, 3, 4)
+    else:
+        raise ValueError(f"unsupported image size: {image_size}")
+
+    attention_ds = []
+    for res in classifier_attention_resolutions.split(","):
+        attention_ds.append(image_size // int(res))
+
+    return EncoderUNetModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=classifier_width,
+        out_channels=1000,
+        num_res_blocks=classifier_depth,
+        attention_resolutions=tuple(attention_ds),
+        channel_mult=channel_mult,
+        use_fp16=classifier_use_fp16,
+        num_head_channels=64,
+        use_scale_shift_norm=classifier_use_scale_shift_norm,
+        resblock_updown=classifier_resblock_updown,
+        pool=classifier_pool,
+    )
+
+
+def sr_model_and_diffusion_defaults():
+    res = model_and_diffusion_defaults()
+    res["large_size"] = 256
+    res["small_size"] = 64
+    arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0]
+    for k in res.copy().keys():
+        if k not in arg_names:
+            del res[k]
+    return res
+
+
+def sr_create_model_and_diffusion(
+    large_size,
+    small_size,
+    class_cond,
+    learn_sigma,
+    num_channels,
+    num_res_blocks,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    attention_resolutions,
+    dropout,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    use_checkpoint,
+    use_scale_shift_norm,
+    resblock_updown,
+    use_fp16,
+):
+    model = sr_create_model(
+        large_size,
+        small_size,
+        num_channels,
+        num_res_blocks,
+        learn_sigma=learn_sigma,
+        class_cond=class_cond,
+        use_checkpoint=use_checkpoint,
+        attention_resolutions=attention_resolutions,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        dropout=dropout,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+    )
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return model, diffusion
+
+
+def sr_create_model(
+    large_size,
+    small_size,
+    num_channels,
+    num_res_blocks,
+    learn_sigma,
+    class_cond,
+    use_checkpoint,
+    attention_resolutions,
+    num_heads,
+    num_head_channels,
+    num_heads_upsample,
+    use_scale_shift_norm,
+    dropout,
+    resblock_updown,
+    use_fp16,
+):
+    _ = small_size  # hack to prevent unused variable
+
+    if large_size == 512:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif large_size == 256:
+        channel_mult = (1, 1, 2, 2, 4, 4)
+    elif large_size == 64:
+        channel_mult = (1, 2, 3, 4)
+    else:
+        raise ValueError(f"unsupported large size: {large_size}")
+
+    attention_ds = []
+    for res in attention_resolutions.split(","):
+        attention_ds.append(large_size // int(res))
+
+    return SuperResModel(
+        image_size=large_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=dropout,
+        channel_mult=channel_mult,
+        num_classes=(NUM_CLASSES if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        num_heads=num_heads,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=num_heads_upsample,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_fp16=use_fp16,
+    )
+
+
+def create_gaussian_diffusion(
+    *,
+    steps=1000,
+    learn_sigma=False,
+    sigma_small=False,
+    noise_schedule="linear",
+    use_kl=False,
+    predict_xstart=False,
+    rescale_timesteps=False,
+    rescale_learned_sigmas=False,
+    timestep_respacing="",
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, steps)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if not timestep_respacing:
+        timestep_respacing = [steps]
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=(
+            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
+        ),
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type,
+        rescale_timesteps=rescale_timesteps,
+    )
+
+
+def add_dict_to_argparser(parser, default_dict):
+    for k, v in default_dict.items():
+        v_type = type(v)
+        if v is None:
+            v_type = str
+        elif isinstance(v, bool):
+            v_type = str2bool
+        parser.add_argument(f"--{k}", default=v, type=v_type)
+
+
+def args_to_dict(args, keys):
+    return {k: getattr(args, k) for k in keys}
+
+def dict_parse(args, keys):
+    return {k: args[k] for k in keys}
+
+
+def str2bool(v):
+    """
+    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("boolean value expected")

guided_diffusion/guided_diffusion/train_util.py

@@ -0,0 +1,301 @@
+import copy
+import functools
+import os
+
+import blobfile as bf
+import torch as th
+import torch.distributed as dist
+from torch.nn.parallel.distributed import DistributedDataParallel as DDP
+from torch.optim import AdamW
+
+from . import dist_util, logger
+from .fp16_util import MixedPrecisionTrainer
+from .nn import update_ema
+from .resample import LossAwareSampler, UniformSampler
+
+# For ImageNet experiments, this was a good default value.
+# We found that the lg_loss_scale quickly climbed to
+# 20-21 within the first ~1K steps of training.
+INITIAL_LOG_LOSS_SCALE = 20.0
+
+
+class TrainLoop:
+    def __init__(
+        self,
+        *,
+        model,
+        diffusion,
+        data,
+        batch_size,
+        microbatch,
+        lr,
+        ema_rate,
+        log_interval,
+        save_interval,
+        resume_checkpoint,
+        use_fp16=False,
+        fp16_scale_growth=1e-3,
+        schedule_sampler=None,
+        weight_decay=0.0,
+        lr_anneal_steps=0,
+    ):
+        self.model = model
+        self.diffusion = diffusion
+        self.data = data
+        self.batch_size = batch_size
+        self.microbatch = microbatch if microbatch > 0 else batch_size
+        self.lr = lr
+        self.ema_rate = (
+            [ema_rate]
+            if isinstance(ema_rate, float)
+            else [float(x) for x in ema_rate.split(",")]
+        )
+        self.log_interval = log_interval
+        self.save_interval = save_interval
+        self.resume_checkpoint = resume_checkpoint
+        self.use_fp16 = use_fp16
+        self.fp16_scale_growth = fp16_scale_growth
+        self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
+        self.weight_decay = weight_decay
+        self.lr_anneal_steps = lr_anneal_steps
+
+        self.step = 0
+        self.resume_step = 0
+        self.global_batch = self.batch_size * dist.get_world_size()
+
+        self.sync_cuda = th.cuda.is_available()
+
+        self._load_and_sync_parameters()
+        self.mp_trainer = MixedPrecisionTrainer(
+            model=self.model,
+            use_fp16=self.use_fp16,
+            fp16_scale_growth=fp16_scale_growth,
+        )
+
+        self.opt = AdamW(
+            self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay
+        )
+        if self.resume_step:
+            self._load_optimizer_state()
+            # Model was resumed, either due to a restart or a checkpoint
+            # being specified at the command line.
+            self.ema_params = [
+                self._load_ema_parameters(rate) for rate in self.ema_rate
+            ]
+        else:
+            self.ema_params = [
+                copy.deepcopy(self.mp_trainer.master_params)
+                for _ in range(len(self.ema_rate))
+            ]
+
+        if th.cuda.is_available():
+            self.use_ddp = True
+            self.ddp_model = DDP(
+                self.model,
+                device_ids=[dist_util.dev()],
+                output_device=dist_util.dev(),
+                broadcast_buffers=False,
+                bucket_cap_mb=128,
+                find_unused_parameters=False,
+            )
+        else:
+            if dist.get_world_size() > 1:
+                logger.warn(
+                    "Distributed training requires CUDA. "
+                    "Gradients will not be synchronized properly!"
+                )
+            self.use_ddp = False
+            self.ddp_model = self.model
+
+    def _load_and_sync_parameters(self):
+        resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+
+        if resume_checkpoint:
+            self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
+            if dist.get_rank() == 0:
+                logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
+                self.model.load_state_dict(
+                    dist_util.load_state_dict(
+                        resume_checkpoint, map_location=dist_util.dev()
+                    )
+                )
+
+        dist_util.sync_params(self.model.parameters())
+
+    def _load_ema_parameters(self, rate):
+        ema_params = copy.deepcopy(self.mp_trainer.master_params)
+
+        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+        ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate)
+        if ema_checkpoint:
+            if dist.get_rank() == 0:
+                logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...")
+                state_dict = dist_util.load_state_dict(
+                    ema_checkpoint, map_location=dist_util.dev()
+                )
+                ema_params = self.mp_trainer.state_dict_to_master_params(state_dict)
+
+        dist_util.sync_params(ema_params)
+        return ema_params
+
+    def _load_optimizer_state(self):
+        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+        opt_checkpoint = bf.join(
+            bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt"
+        )
+        if bf.exists(opt_checkpoint):
+            logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}")
+            state_dict = dist_util.load_state_dict(
+                opt_checkpoint, map_location=dist_util.dev()
+            )
+            self.opt.load_state_dict(state_dict)
+
+    def run_loop(self):
+        while (
+            not self.lr_anneal_steps
+            or self.step + self.resume_step < self.lr_anneal_steps
+        ):
+            batch, cond = next(self.data)
+            self.run_step(batch, cond)
+            if self.step % self.log_interval == 0:
+                logger.dumpkvs()
+            if self.step % self.save_interval == 0:
+                self.save()
+                # Run for a finite amount of time in integration tests.
+                if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
+                    return
+            self.step += 1
+        # Save the last checkpoint if it wasn't already saved.
+        if (self.step - 1) % self.save_interval != 0:
+            self.save()
+
+    def run_step(self, batch, cond):
+        self.forward_backward(batch, cond)
+        took_step = self.mp_trainer.optimize(self.opt)
+        if took_step:
+            self._update_ema()
+        self._anneal_lr()
+        self.log_step()
+
+    def forward_backward(self, batch, cond):
+        self.mp_trainer.zero_grad()
+        for i in range(0, batch.shape[0], self.microbatch):
+            micro = batch[i : i + self.microbatch].to(dist_util.dev())
+            micro_cond = {
+                k: v[i : i + self.microbatch].to(dist_util.dev())
+                for k, v in cond.items()
+            }
+            last_batch = (i + self.microbatch) >= batch.shape[0]
+            t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())
+
+            compute_losses = functools.partial(
+                self.diffusion.training_losses,
+                self.ddp_model,
+                micro,
+                t,
+                model_kwargs=micro_cond,
+            )
+
+            if last_batch or not self.use_ddp:
+                losses = compute_losses()
+            else:
+                with self.ddp_model.no_sync():
+                    losses = compute_losses()
+
+            if isinstance(self.schedule_sampler, LossAwareSampler):
+                self.schedule_sampler.update_with_local_losses(
+                    t, losses["loss"].detach()
+                )
+
+            loss = (losses["loss"] * weights).mean()
+            log_loss_dict(
+                self.diffusion, t, {k: v * weights for k, v in losses.items()}
+            )
+            self.mp_trainer.backward(loss)
+
+    def _update_ema(self):
+        for rate, params in zip(self.ema_rate, self.ema_params):
+            update_ema(params, self.mp_trainer.master_params, rate=rate)
+
+    def _anneal_lr(self):
+        if not self.lr_anneal_steps:
+            return
+        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
+        lr = self.lr * (1 - frac_done)
+        for param_group in self.opt.param_groups:
+            param_group["lr"] = lr
+
+    def log_step(self):
+        logger.logkv("step", self.step + self.resume_step)
+        logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)
+
+    def save(self):
+        def save_checkpoint(rate, params):
+            state_dict = self.mp_trainer.master_params_to_state_dict(params)
+            if dist.get_rank() == 0:
+                logger.log(f"saving model {rate}...")
+                if not rate:
+                    filename = f"model{(self.step+self.resume_step):06d}.pt"
+                else:
+                    filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt"
+                with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f:
+                    th.save(state_dict, f)
+
+        save_checkpoint(0, self.mp_trainer.master_params)
+        for rate, params in zip(self.ema_rate, self.ema_params):
+            save_checkpoint(rate, params)
+
+        if dist.get_rank() == 0:
+            with bf.BlobFile(
+                bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"),
+                "wb",
+            ) as f:
+                th.save(self.opt.state_dict(), f)
+
+        dist.barrier()
+
+
+def parse_resume_step_from_filename(filename):
+    """
+    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
+    checkpoint's number of steps.
+    """
+    split = filename.split("model")
+    if len(split) < 2:
+        return 0
+    split1 = split[-1].split(".")[0]
+    try:
+        return int(split1)
+    except ValueError:
+        return 0
+
+
+def get_blob_logdir():
+    # You can change this to be a separate path to save checkpoints to
+    # a blobstore or some external drive.
+    return logger.get_dir()
+
+
+def find_resume_checkpoint():
+    # On your infrastructure, you may want to override this to automatically
+    # discover the latest checkpoint on your blob storage, etc.
+    return None
+
+
+def find_ema_checkpoint(main_checkpoint, step, rate):
+    if main_checkpoint is None:
+        return None
+    filename = f"ema_{rate}_{(step):06d}.pt"
+    path = bf.join(bf.dirname(main_checkpoint), filename)
+    if bf.exists(path):
+        return path
+    return None
+
+
+def log_loss_dict(diffusion, ts, losses):
+    for key, values in losses.items():
+        logger.logkv_mean(key, values.mean().item())
+        # Log the quantiles (four quartiles, in particular).
+        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
+            quartile = int(4 * sub_t / diffusion.num_timesteps)
+            logger.logkv_mean(f"{key}_q{quartile}", sub_loss)

guided_diffusion/guided_diffusion/unet.py

@@ -0,0 +1,897 @@
+from abc import abstractmethod
+
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .fp16_util import convert_module_to_f16, convert_module_to_f32
+from .nn import (
+    checkpoint,
+    conv_nd,
+    linear,
+    avg_pool_nd,
+    zero_module,
+    normalization,
+    timestep_embedding,
+)
+
+
+class AttentionPool2d(nn.Module):
+    """
+    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+    """
+
+    def __init__(
+        self,
+        spacial_dim: int,
+        embed_dim: int,
+        num_heads_channels: int,
+        output_dim: int = None,
+    ):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(
+            th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5
+        )
+        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+        self.num_heads = embed_dim // num_heads_channels
+        self.attention = QKVAttention(self.num_heads)
+
+    def forward(self, x):
+        b, c, *_spatial = x.shape
+        x = x.reshape(b, c, -1)  # NC(HW)
+        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)
+        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)
+        x = self.qkv_proj(x)
+        x = self.attention(x)
+        x = self.c_proj(x)
+        return x[:, :, 0]
+
+
+class TimestepBlock(nn.Module):
+    """
+    Any module where forward() takes timestep embeddings as a second argument.
+    """
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    """
+    A sequential module that passes timestep embeddings to the children that
+    support it as an extra input.
+    """
+
+    def forward(self, x, emb):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=1
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+    """
+    A residual block that can optionally change the number of channels.
+
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        return checkpoint(
+            self._forward, (x, emb), self.parameters(), self.use_checkpoint
+        )
+
+    def _forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = th.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        use_checkpoint=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.use_checkpoint = use_checkpoint
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, channels * 3, 1)
+        if use_new_attention_order:
+            # split qkv before split heads
+            self.attention = QKVAttention(self.num_heads)
+        else:
+            # split heads before split qkv
+            self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+    def forward(self, x):
+        return checkpoint(self._forward, (x,), self.parameters(), True)
+
+    def _forward(self, x):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+def count_flops_attn(model, _x, y):
+    """
+    A counter for the `thop` package to count the operations in an
+    attention operation.
+    Meant to be used like:
+        macs, params = thop.profile(
+            model,
+            inputs=(inputs, timestamps),
+            custom_ops={QKVAttention: QKVAttention.count_flops},
+        )
+    """
+    b, c, *spatial = y[0].shape
+    num_spatial = int(np.prod(spatial))
+    # We perform two matmuls with the same number of ops.
+    # The first computes the weight matrix, the second computes
+    # the combination of the value vectors.
+    matmul_ops = 2 * b * (num_spatial ** 2) * c
+    model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+    """
+    A module which performs QKV attention and splits in a different order.
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.chunk(3, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class UNetModel(nn.Module):
+    """
+    The full UNet model with attention and timestep embedding.
+
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.dtype = th.float16 if use_fp16 else th.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if self.num_classes is not None:
+            self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+
+        ch = input_ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(model_channels * mult),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(model_channels * mult)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads_upsample,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    layers.append(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
+        )
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps, y=None):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        # assert (y is not None) == (
+        #     self.num_classes is not None
+        # ), "must specify y if and only if the model is class-conditional"
+        
+        hs = []
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        if self.num_classes is not None:
+            if y is None:
+                pass
+            else:
+                assert y.shape == (x.shape[0],)
+                emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            hs.append(h)
+        h = self.middle_block(h, emb)
+        for module in self.output_blocks:
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(h, emb)
+        h = h.type(x.dtype)
+        return self.out(h)
+
+
+class SuperResModel(UNetModel):
+    """
+    A UNetModel that performs super-resolution.
+
+    Expects an extra kwarg `low_res` to condition on a low-resolution image.
+    """
+
+    def __init__(self, image_size, in_channels, *args, **kwargs):
+        super().__init__(image_size, in_channels * 2, *args, **kwargs)
+
+    def forward(self, x, timesteps, low_res=None, **kwargs):
+        _, _, new_height, new_width = x.shape
+        upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
+        x = th.cat([x, upsampled], dim=1)
+        return super().forward(x, timesteps, **kwargs)
+
+
+class EncoderUNetModel(nn.Module):
+    """
+    The half UNet model with attention and timestep embedding.
+
+    For usage, see UNet.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        pool="adaptive",
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.dtype = th.float16 if use_fp16 else th.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+        self.pool = pool
+        if pool == "adaptive":
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                nn.AdaptiveAvgPool2d((1, 1)),
+                zero_module(conv_nd(dims, ch, out_channels, 1)),
+                nn.Flatten(),
+            )
+        elif pool == "attention":
+            assert num_head_channels != -1
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                AttentionPool2d(
+                    (image_size // ds), ch, num_head_channels, out_channels
+                ),
+            )
+        elif pool == "spatial":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                nn.ReLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        elif pool == "spatial_v2":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                normalization(2048),
+                nn.SiLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        else:
+            raise NotImplementedError(f"Unexpected {pool} pooling")
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :return: an [N x K] Tensor of outputs.
+        """
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        results = []
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            if self.pool.startswith("spatial"):
+                results.append(h.type(x.dtype).mean(dim=(2, 3)))
+        h = self.middle_block(h, emb)
+        if self.pool.startswith("spatial"):
+            results.append(h.type(x.dtype).mean(dim=(2, 3)))
+            h = th.cat(results, axis=-1)
+            return self.out(h)
+        else:
+            h = h.type(x.dtype)
+            return self.out(h)

guided_diffusion/model-card.md

@@ -0,0 +1,59 @@
+# Overview
+
+These are diffusion models and noised image classifiers described in the paper [Diffusion Models Beat GANs on Image Synthesis](https://arxiv.org/abs/2105.05233).
+Included in this release are the following models:
+
+ * Noisy ImageNet classifiers at resolutions 64x64, 128x128, 256x256, 512x512
+ * A class-unconditional ImageNet diffusion model at resolution 256x256
+ * Class conditional ImageNet diffusion models at 64x64, 128x128, 256x256, 512x512 resolutions
+ * Class-conditional ImageNet upsampling diffusion models: 64x64->256x256, 128x128->512x512
+ * Diffusion models trained on three LSUN classes at 256x256 resolution: cat, horse, bedroom
+
+# Datasets
+
+All of the models we are releasing were either trained on the [ILSVRC 2012 subset of ImageNet](http://www.image-net.org/challenges/LSVRC/2012/) or on single classes of [LSUN](https://arxiv.org/abs/1506.03365).
+Here, we describe characteristics of these datasets which impact model behavior:
+
+**LSUN**: This dataset was collected in 2015 using a combination of human labeling (from Amazon Mechanical Turk) and automated data labeling.
+ * Each of the three classes we consider contain over a million images.
+ * The dataset creators found that the label accuracy was roughly 90% across the entire LSUN dataset when measured by trained experts.
+ * Images are scraped from the internet, and LSUN cat images in particular tend to often follow a “meme” format.
+ * We found that there are occasionally humans in these photos, including faces, especially within the cat class.  
+
+**ILSVRC 2012 subset of ImageNet**: This dataset was curated in 2012 and consists of roughly one million images, each belonging to one of 1000 classes.
+ * A large portion of the classes in this dataset are animals, plants, and other naturally-occurring objects.
+ * Many images contain humans, although usually these humans aren’t reflected by the class label (e.g. the class “Tench, tinca tinca” contains many photos of people holding fish).
+
+# Performance
+
+These models are intended to generate samples consistent with their training distributions.
+This has been measured in terms of FID, Precision, and Recall.
+These metrics all rely on the representations of a [pre-trained Inception-V3 model](https://arxiv.org/abs/1512.00567),
+which was trained on ImageNet, and so is likely to focus more on the ImageNet classes (such as animals) than on other visual features (such as human faces).
+
+Qualitatively, the samples produced by these models often look highly realistic, especially when a diffusion model is combined with a noisy classifier.
+
+# Intended Use
+
+These models are intended to be used for research purposes only.
+In particular, they can be used as a baseline for generative modeling research, or as a starting point to build off of for such research.
+
+These models are not intended to be commercially deployed.
+Additionally, they are not intended to be used to create propaganda or offensive imagery.
+
+Before releasing these models, we probed their ability to ease the creation of targeted imagery, since doing so could be potentially harmful.
+We did this either by fine-tuning our ImageNet models on a target LSUN class, or through classifier guidance with publicly available [CLIP models](https://github.com/openai/CLIP).
+ * To probe fine-tuning capabilities, we restricted our compute budget to roughly $100 and tried both standard fine-tuning,
+and a diffusion-specific approach where we train a specialized classifier for the LSUN class. The resulting FIDs were significantly worse than publicly available GAN models, indicating that fine-tuning an ImageNet diffusion model does not significantly lower the cost of image generation.
+ * To probe guidance with CLIP, we tried two approaches for using pre-trained CLIP models for classifier guidance. Either we fed the noised image to CLIP directly and used its gradients, or we fed the diffusion model's denoised prediction to the CLIP model and differentiated through the whole process. In both cases, we found that it was difficult to recover information from the CLIP model, indicating that these diffusion models are unlikely to make it significantly easier to extract knowledge from CLIP compared to existing GAN models.
+
+# Limitations
+
+These models sometimes produce highly unrealistic outputs, particularly when generating images containing human faces.
+This may stem from ImageNet's emphasis on non-human objects.
+
+While classifier guidance can improve sample quality, it reduces diversity, resulting in some modes of the data distribution being underrepresented.
+This can potentially amplify existing biases in the training dataset such as gender and racial biases.
+
+Because ImageNet and LSUN contain images from the internet, they include photos of real people, and the model may have memorized some of the information contained in these photos.
+However, these images are already publicly available, and existing generative models trained on ImageNet have not demonstrated significant leakage of this information.

guided_diffusion/setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup
+
+setup(
+    name="guided-diffusion",
+    py_modules=["guided_diffusion"],
+    install_requires=["blobfile>=1.0.5", "torch", "tqdm"],
+)

networks/distill_model.py

@@ -0,0 +1,76 @@
+import torch.nn as nn
+import torchvision.models as TVM
+from collections import OrderedDict
+from guided_diffusion.guided_diffusion.fp16_util import convert_module_to_f16
+
+import torch
+
+# ------------------------------------------------------------------------------
+# Model Definition
+class DistilDIRE(torch.nn.Module):
+    def __init__(self, device):
+        super(DistilDIRE, self).__init__()
+    
+        # define models
+        student = TVM.resnet50()
+        self.student_backbone = nn.Sequential(OrderedDict([*(list(student.named_children())[:-2])])) # drop last layer which is classifier
+        # extract last classifier head from teacher resnet 
+        self.student_head = nn.Sequential(nn.AdaptiveAvgPool2d(1),
+                                          nn.Flatten(),
+                                          nn.Linear(2048, 1))
+        self.student_backbone.conv1 = nn.Conv2d(6, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.device = device
+    
+    def convert_to_fp16_student(self):
+        self.student_backbone.apply(convert_module_to_f16)
+    
+
+    def forward(self, img, eps):
+        img = img.to(self.device)
+        eps = eps.to(self.device)
+    
+        # concat image and noise
+        img_tens = torch.cat([img, eps], dim=1)
+        
+        feature = self.student_backbone(img_tens) 
+        logit = self.student_head(feature)
+        return {'logit':logit, 'feature':feature}
+    
+
+# ------------------------------------------------------------------------------
+class DistilDIREOnlyEPS(DistilDIRE):
+    def __init__(self, device):
+        super(DistilDIREOnlyEPS, self).__init__(device)
+        self.student_backbone.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.student_head = nn.Sequential(nn.AdaptiveAvgPool2d(1),
+                                          nn.Flatten(),
+                                          nn.Linear(2048, 1))
+        
+    def forward(self, eps):
+        eps = eps.to(self.device)
+        feature = self.student_backbone(eps)
+        logit = self.student_head(feature)
+        return {'logit':logit, 'feature':feature}
+
+
+# ------------------------------------------------------------------------------
+class DIRE(torch.nn.Module):
+    def __init__(self, device):
+        super(DIRE, self).__init__()
+    
+        # define models
+        student = TVM.resnet50()
+        self.student_backbone = nn.Sequential(OrderedDict([*(list(student.named_children())[:-2])])) # drop last layer which is classifier
+        # extract last classifier head from teacher resnet 
+        self.student_head = nn.Sequential(nn.AdaptiveAvgPool2d(1),
+                                          nn.Flatten(),
+                                          nn.Linear(2048, 1))
+        self.device = device
+        
+        
+    def forward(self, dire):
+        dire = dire.to(self.device)
+    
+        feature = self.student_backbone(dire) 
+        logit = self.student_head(feature)
+        return {'logit':logit, 'feature':feature}

requirements.txt

@@ -0,0 +1,21 @@
+# conda create -n dire python=3.9
+# pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+einops
+imageio
+ipympl
+cmake==3.28.3
+dlib==19.24.2
+decord
+face_recognition
+matplotlib
+numpy
+opencv-python==4.9.0.80
+pandas
+scikit-learn
+tensorboard
+tensorboardX
+sanic
+wandb
+tqdm
+blobfile>=1.0.5
+easydict

server.py

@@ -0,0 +1,55 @@
+# Two server routes that OctoAI containers should have:
+# a route for inference requests (e.g. ”/predict”). This route for inference requests must receive JSON inputs and JSON outputs.
+# a route for health checks (e.g. ”/healthcheck”).
+# Number of workers (not required). Typical best practice is to make this number some function of the # of CPU cores that the server has access to and should use.
+
+"""HTTP Inference serving interface using sanic."""
+import os
+
+from custommodel import CustomModel
+from sanic import Request, Sanic, response
+
+_DEFAULT_PORT = 8000
+"""Default port to serve inference on."""
+
+# Load and initialize the model on startup globally, so it can be reused.
+model_instance = CustomModel(ckpt='models/distil-240709-tot-model_epoch_20.pth')
+"""Global instance of the model to serve."""
+
+server = Sanic("server")
+"""Global instance of the web server."""
+
+
+@server.route("/healthcheck", methods=["GET"])
+def healthcheck(_: Request) -> response.JSONResponse:
+    """Responds to healthcheck requests.
+
+    :param request: the incoming healthcheck request.
+    :return: json responding to the healthcheck.
+    """
+    return response.json({"healthy": "yes"})
+
+
+@server.route("/reconstruction/predict", methods=["POST"])
+def predict(request: Request) -> response.JSONResponse:
+    """Responds to inference/prediction requests.
+
+    :param request: the incoming request containing inputs for the model.
+    :return: json containing the inference results.
+    """
+
+    try:
+        inputs = request.json
+        output = model_instance.predict(inputs)
+        return response.json(output)
+    except Exception as e:
+        return response.json({'error': str(e)}, status=500)
+
+
+def main():
+    """Entry point for the server."""
+    port = int(os.environ.get("SERVING_PORT", _DEFAULT_PORT))
+    server.run(host="0.0.0.0", port=port, workers=1)
+
+if __name__ == "__main__":
+    main()

test.py

@@ -0,0 +1,22 @@
+import wandb
+from utils.config import cfg
+
+    
+def main(run, cfg):
+    from utils.trainer import Trainer
+    from torch.utils.data import DataLoader
+    from dataset import TMEPSOnlyDataset, TMIMGOnlyDataset
+    print(cfg.dataset_test_root)
+    # dataset =TMEPSOnlyDataset(cfg.dataset_test_root, False)
+    dataset = TMIMGOnlyDataset(cfg.dataset_test_root, istrain=False)
+    dataloader = DataLoader(dataset, 
+                            batch_size=1, 
+                            shuffle=True, num_workers=2)
+    trainer = Trainer(cfg, dataloader, dataloader, run, 0, False, 1)
+    assert len(cfg.pretrained_weights) != 0, "Give proper checkpoint path"
+    trainer.load_networks(cfg.pretrained_weights)
+    trainer.validate(False, save=True, save_name=f"{cfg.root_dir}/{cfg.datasets_test}_{cfg.pretrained_weights.split('/')[-1]}_results.txt")
+
+if __name__ == "__main__":
+    main(None, cfg)
+    

train.py

@@ -0,0 +1,59 @@
+
+import os 
+import torch
+
+
+def main(run, cfg):
+    from torch.utils.data.distributed import DistributedSampler
+    from utils.trainer import Trainer
+    
+    if cfg.reproduce_dire:
+        dataset = TMDireDataset(cfg.dataset_root)
+        val_dataset = TMDireDataset(cfg.dataset_test_root)
+    elif cfg.only_eps:
+        dataset = TMEPSOnlyDataset(cfg.dataset_root)
+        val_dataset = TMEPSOnlyDataset(cfg.dataset_root)
+    
+    elif cfg.only_img:
+        dataset = TMIMGOnlyDataset(cfg.dataset_root, istrain=True)
+        val_dataset = TMIMGOnlyDataset(cfg.dataset_test_root, istrain=False)
+    else:
+        dataset= TMDistilDireDataset(cfg.dataset_root)
+        val_dataset = TMDistilDireDataset(cfg.dataset_test_root)
+    sampler = DistributedSampler(dataset)
+    val_samlper = DistributedSampler(val_dataset)
+    dataloader = DataLoader(dataset, 
+                            batch_size=cfg.batch_size, 
+                            sampler=sampler,
+                            num_workers=2)
+    val_loader = DataLoader(val_dataset, 
+                            batch_size=cfg.batch_size, 
+                            sampler=val_samlper,
+                            num_workers=2)
+    trainer = Trainer(cfg, dataloader, val_loader, run, local_rank, True, world_size, cfg.kd)
+    if cfg.pretrained_weights:
+        trainer.load_networks(cfg.pretrained_weights)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    
+    import torch.distributed as dist
+    import os 
+    import wandb
+    
+    from torch.utils.data import DataLoader
+    from dataset import TMDistilDireDataset, TMDireDataset, TMEPSOnlyDataset, TMIMGOnlyDataset
+    
+    dist.init_process_group(backend='nccl', init_method='env://')
+    local_rank = int(os.environ['LOCAL_RANK']) 
+    world_size = int(os.environ['WORLD_SIZE'])
+    torch.cuda.set_device(local_rank)
+    dist.barrier()
+    from utils.config import cfg
+    run = None
+    if local_rank == 0:
+        run = wandb.init(project=f'dire-distill-truemedia', config=cfg, dir=cfg.exp_dir) 
+    main(run, cfg)
+    
+

train.sh

@@ -0,0 +1,3 @@
+#! /bin/bash
+
+torchrun --standalone --nproc_per_node 1 -m train --batch 64 --exp_name 240709-tot-finetune-from-scratch-fp16-distildire-w-prp-pos2 --datasets truemedia-total --datasets_test truemedia-external --epoch 50 --lr 1e-4 --pretrained_weights /home/ubuntu/y1/DistilDIRE/experiments/240709-tot-finetune-from-scratch-fp16-distildire-w-prp/ckpt/model_epoch_4.pth 

utils/config.py

@@ -0,0 +1,179 @@
+import argparse
+import os
+import sys
+from abc import ABC
+from typing import Type
+
+
+class DefaultConfigs(ABC):
+    ####### base setting ######
+    gpus = [0]
+    seed = 3407
+    arch = "resnet50"
+    datasets = [""]
+    datasets_test = [""]
+    class_bal = False
+    batch_size = 256
+    val_every = 1
+    loadSize = 256
+    cropSize = 224
+    epoch = "latest"
+    num_workers = 2
+    isTrain = True
+
+    ####### train setting ######
+    warmup = False
+    warmup_epoch = 3
+    earlystop = True
+    earlystop_epoch = 5
+    optim = "adam"
+    new_optim = False
+    loss_freq = 400
+    save_latest_freq = 2000
+    save_epoch_freq = 20
+    continue_train = False
+    epoch_count = 1
+    last_epoch = -1
+    nepoch = 20
+    beta1 = 0.9
+    lr = 0.00001
+    init_type = "normal"
+    init_gain = 0.02
+    pretrained = True
+
+    # paths information
+    root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    dataset_root = os.path.join(root_dir, "datasets")
+    dataset_test_root = os.path.join(root_dir, "datasets")
+    exp_root = os.path.join(root_dir, "experiments")
+    _exp_name = ""
+    exp_dir = ""
+    ckpt_dir = ""
+    logs_path = ""
+    ckpt_path = ""
+    pretrained_weights = ""
+    kd = True
+    kd_weight = 1.
+    reproduce_dire = False
+    only_eps = False
+    only_img = False
+
+    @property
+    def exp_name(self):
+        return self._exp_name
+
+    @exp_name.setter
+    def exp_name(self, value: str):
+        self._exp_name = value
+        self.exp_dir: str = os.path.join(self.exp_root, self.exp_name)
+        self.ckpt_dir: str = os.path.join(self.exp_dir, "ckpt")
+        self.logs_path: str = os.path.join(self.exp_dir, "logs.txt")
+        if self.isTrain:
+            os.makedirs(self.exp_dir, exist_ok=True)
+            os.makedirs(self.ckpt_dir, exist_ok=True)
+
+    def to_dict(self):
+        dic = {}
+        for fieldkey in dir(self):
+            fieldvalue = getattr(self, fieldkey)
+            if not fieldkey.startswith("__") and not callable(fieldvalue) and not fieldkey.startswith("_"):
+                dic[fieldkey] = fieldvalue
+        return dic
+
+
+def args_list2dict(arg_list: list):
+    assert len(arg_list) % 2 == 0, f"Override list has odd length: {arg_list}; it must be a list of pairs"
+    return dict(zip(arg_list[::2], arg_list[1::2]))
+
+
+def str2bool(v: str) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif v.lower() in ("true", "yes", "on", "y", "t", "1"):
+        return True
+    elif v.lower() in ("false", "no", "off", "n", "f", "0"):
+        return False
+    else:
+        return bool(v)
+
+
+def str2list(v: str, element_type=None) -> list:
+    if not isinstance(v, (list, tuple, set)):
+        v = v.lstrip("[").rstrip("]")
+        v = v.split(",")
+        v = list(map(str.strip, v))
+        if element_type is not None:
+            v = list(map(element_type, v))
+    return v
+
+
+CONFIGCLASS = Type[DefaultConfigs]
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpus", default="0", type=str)
+parser.add_argument("--batch", default=256, type=int)
+parser.add_argument("--epoch", default="100", type=int)
+parser.add_argument("--exp_name", default="", type=str)
+parser.add_argument("--datasets", default="", type=str)
+parser.add_argument("--dataset_root", default="", type=str)
+parser.add_argument("--datasets_test", default="", type=str)
+parser.add_argument("--dataset_test_root", default="", type=str)
+parser.add_argument("--pretrained_weights", default="", type=str)
+parser.add_argument("--lr", default=0.00001, type=float)
+parser.add_argument("--kd_weight", default=1., type=float)
+parser.add_argument("--test", default=False, type=str2bool)
+parser.add_argument("--reproduce_dire", default=False, type=str2bool)
+parser.add_argument("--only_eps", default=False, type=str2bool)
+parser.add_argument("--only_img", default=False, type=str2bool)
+parser.add_argument("--kd", default=True, type=str2bool)
+parser.add_argument("opts", default=[], nargs=argparse.REMAINDER)
+args = parser.parse_args()
+
+if os.path.exists(os.path.join(DefaultConfigs.exp_root, args.exp_name, "config.py")):
+    sys.path.insert(0, os.path.join(DefaultConfigs.exp_root, args.exp_name))
+    from config import cfg
+
+    cfg: CONFIGCLASS
+else:
+    cfg = DefaultConfigs()
+
+if args.opts:
+    opts = args_list2dict(args.opts)
+    for k, v in opts.items():
+        if not hasattr(cfg, k):
+            raise ValueError(f"Unrecognized option: {k}")
+        original_type = type(getattr(cfg, k))
+        if original_type == bool:
+            setattr(cfg, k, str2bool(v))
+        elif original_type in (list, tuple, set):
+            setattr(cfg, k, str2list(v, type(getattr(cfg, k)[0])))
+        else:
+            setattr(cfg, k, original_type(v))
+
+#cfg.gpus: list = args.gpus
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus#", ".join([str(gpu) for gpu in cfg.gpus])
+if args.test:
+    cfg.isTrain = False
+cfg.exp_name = args.exp_name
+cfg.batch_size = args.batch
+cfg.datasets = args.datasets
+cfg.datasets_test = args.datasets_test if args.datasets_test else args.datasets
+cfg.pretrained_weights = args.pretrained_weights
+cfg.lr = args.lr
+cfg.nepoch = args.epoch
+
+root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+cfg.dataset_root = os.path.join(root_dir, 'datasets', cfg.datasets)
+if args.dataset_root != "":
+    cfg.dataset_root = args.dataset_root
+cfg.dataset_test_root = os.path.join(root_dir, 'datasets', cfg.datasets_test)
+if args.dataset_test_root != "":
+    cfg.dataset_test_root = args.dataset_test_root
+cfg.kd = args.kd
+cfg.reproduce_dire = args.reproduce_dire
+cfg.only_eps = args.only_eps
+cfg.only_img = args.only_img
+cfg.kd_weight = args.kd_weight
+
+# if isinstance(cfg.datasets, str):
+#     cfg.datasets = cfg.datasets.split(",")

utils/sanic_utils.py

@@ -0,0 +1,88 @@
+from decord import VideoReader,cpu
+import os 
+import numpy as np
+import cv2 
+import dlib
+import face_recognition
+from tqdm import tqdm
+
+def real_or_fake(prediction):
+    return {0: "REAL", 1: "FAKE"}[prediction ^ 1]
+
+
+def real_or_fake_thres(probability, threshold=0.2):
+    return "FAKE" if probability >= threshold else "REAL"
+
+
+def face_rec(frames, p=None, klass=None):
+    temp_face = np.zeros((len(frames), 224, 224, 3), dtype=np.uint8)
+    count = 0
+    mod = "cnn" if dlib.DLIB_USE_CUDA else "hog"
+
+    for _, frame in tqdm(enumerate(frames), total=len(frames)):
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        face_locations = face_recognition.face_locations(
+            frame, number_of_times_to_upsample=0, model=mod
+        )
+
+        for face_location in face_locations:
+            if count < len(frames):
+                top, right, bottom, left = face_location
+                face_image = frame[top:bottom, left:right]
+                face_image = cv2.resize(
+                    face_image, (224, 224), interpolation=cv2.INTER_AREA
+                )
+                face_image = cv2.cvtColor(face_image, cv2.COLOR_BGR2RGB)
+
+                temp_face[count] = face_image
+                count += 1
+            else:
+                break
+
+    return ([], 0) if count == 0 else (temp_face[:count], count)
+
+
+def extract_frames(video_file, frames_nums=15):
+    vr = VideoReader(video_file, ctx=cpu(0))
+    # Calculate the step size between frames
+    step_size = max(1, len(vr) // frames_nums)
+    return vr.get_batch(
+        list(range(0, len(vr), step_size))[:frames_nums]
+    ).asnumpy()  # seek frames with step_size
+
+def is_video(vid):
+    return os.path.isfile(vid) and vid.endswith(
+        tuple([".avi", ".mp4", ".mpg", ".mpeg", ".mov"])
+    )
+
+def is_image(img):
+    return os.path.isfile(img) and img.endswith(
+        tuple([".jpg", ".jpeg", ".png", ".webp"])
+    )
+
+def set_result():
+    return {
+        "video": {
+            "name": [],
+            "pred": [],
+            "klass": [],
+            "pred_label": [],
+            "correct_label": [],
+        }
+    }
+
+def store_result(
+    result, filename, y, y_val, klass, correct_label=None, compression=None
+):
+    result["video"]["name"].append(filename)
+    result["video"]["pred"].append(y_val)
+    result["video"]["klass"].append(klass.lower())
+    result["video"]["pred_label"].append(real_or_fake(y))
+
+    if correct_label is not None:
+        result["video"]["correct_label"].append(correct_label)
+
+    if compression is not None:
+        result["video"]["compression"].append(compression)
+
+    return result

utils/trainer.py

@@ -0,0 +1,329 @@
+import os
+import torch
+import torch.nn as nn
+
+import torchvision.models as TVM
+from collections import OrderedDict
+from sklearn.metrics import accuracy_score, average_precision_score, precision_score
+from torchvision.io import decode_jpeg, encode_jpeg
+import torch.distributed as dist
+from tqdm.auto import tqdm
+import numpy as np
+
+from utils.config import CONFIGCLASS
+from networks.distill_model import DistilDIRE, DIRE, DistilDIREOnlyEPS
+from utils.warmup import GradualWarmupScheduler
+import os.path as osp
+from guided_diffusion.compute_dire_eps import dire_get_first_step_noise, create_dicts_for_static_init
+from guided_diffusion.guided_diffusion.script_util import (
+                model_and_diffusion_defaults,
+                create_model_and_diffusion,
+                add_dict_to_argparser,
+                dict_parse,
+                args_to_dict,
+)
+
+class BaseModel(nn.Module):
+    def __init__(self, cfg: CONFIGCLASS):
+        super().__init__()
+        self.cfg = cfg
+        self.total_steps = 0
+        self.isTrain = cfg.isTrain
+        self.save_dir = cfg.ckpt_dir
+        self.nepoch = cfg.nepoch
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.student: nn.Module
+        self.optimizer: torch.optim.Optimizer
+
+    def save_networks(self, epoch: int):
+        save_filename = f"model_epoch_{epoch}.pth"
+        save_path = os.path.join(self.save_dir, save_filename)
+
+        # serialize model and optimizer to dict
+        state_dict = {
+            "model": self.student.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "total_steps": self.total_steps,
+        }
+
+        torch.save(state_dict, save_path)
+
+    # load models from the disk
+    def load_networks(self, epoch: int):
+        load_filename = f"model_epoch_{epoch}.pth"
+        load_path = os.path.join(self.save_dir, load_filename)
+
+        print(f"loading the model from {load_path}")
+        # if you are using PyTorch newer than 0.4 (e.g., built from
+        # GitHub source), you can remove str() on self.device
+        state_dict = torch.load(load_path, map_location=self.device)["model"]
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        if hasattr(state_dict, "_metadata"):
+            del state_dict._metadata
+
+        self.student.load_state_dict(state_dict)
+        self.total_steps = state_dict["total_steps"]
+
+        if self.isTrain and not self.cfg.new_optim:
+            self.optimizer.load_state_dict(state_dict["optimizer"])
+            # move optimizer state to GPU
+            for state in self.optimizer.state.values():
+                for k, v in state.items():
+                    if torch.is_tensor(v):
+                        state[k] = v.to(self.device)
+
+            for g in self.optimizer.param_groups:
+                g["lr"] = self.cfg.lr
+
+    def eval(self):
+        self.student.eval()
+
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+
+
+class Trainer(BaseModel):
+    def name(self):
+        return "DistilDIRE Trainer"
+
+    def __init__(self, cfg: CONFIGCLASS, train_loader, val_loader, run, rank=0, distributed=True, world_size=1, kd=True):
+        super().__init__(cfg)
+        self.arch = cfg.arch
+        self.reproduce_dire = cfg.reproduce_dire
+        self.only_eps = cfg.only_eps
+        self.only_img = cfg.only_img
+        self.test_name = osp.basename(cfg.dataset_test_root)
+        self.rank = rank
+        self.device = torch.device(f"cuda") 
+        self.distributed = distributed
+        self.world_size = world_size
+        self.kd = kd
+        self.kd_weight = cfg.kd_weight
+        
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        
+        self.val_every = cfg.val_every
+        self.cur_epoch = 0
+        
+        # wandb logger (pass if None)
+        self.run = run
+        self.adm = None
+        if self.reproduce_dire:
+            self.student = DIRE(self.device).to(self.device)
+        else: 
+            if self.only_eps or self.only_img:
+                self.student = DistilDIREOnlyEPS(self.device).to(self.device)
+                if self.only_img:
+                    adm_args = create_dicts_for_static_init()
+                    adm_args['timestep_respacing'] = 'ddim20'
+                    adm_model, diffusion = create_model_and_diffusion(**dict_parse(adm_args, model_and_diffusion_defaults().keys()))
+                    adm_model.load_state_dict(torch.load(adm_args['model_path'], map_location="cpu"))
+                    print("ADM model loaded...")
+                    self.adm = adm_model
+                    self.adm.convert_to_fp16()
+                    self.adm.to(self.device)
+                    self.adm.eval()
+                    self.diffusion = diffusion
+                    self.adm_args = adm_args
+            else:
+                self.student = DistilDIRE(self.device).to(self.device)
+            # self.student.convert_to_fp16_student()
+            __backbone = TVM.resnet50(weights=TVM.ResNet50_Weights.DEFAULT)
+            self.teacher = nn.Sequential(OrderedDict([*(list(__backbone.named_children())[:-2])])) # drop last layer which is classifier
+            self.teacher.eval().to(self.device)
+            # Freeze teacher model
+            for param in self.teacher.parameters():
+                param.requires_grad = False
+            self.kd_criterion = nn.MSELoss(reduction='mean')
+        
+        self.cls_criterion = nn.BCEWithLogitsLoss(reduction='mean', pos_weight=torch.tensor(0.3))
+        
+        # initialize optimizers
+        if cfg.optim == "adam":
+            self.optimizer = torch.optim.Adam(self.student.parameters(), lr=cfg.lr, betas=(cfg.beta1, 0.999))
+        elif cfg.optim == "sgd":
+            self.optimizer = torch.optim.SGD(self.student.parameters(), lr=cfg.lr, momentum=0.9, weight_decay=5e-4)
+        else:
+            raise ValueError("optim should be [adam, sgd]")
+        
+        if self.distributed:
+            from torch.nn.parallel import DistributedDataParallel as DDP
+            
+            self.student = DDP(self.student, device_ids=[rank], output_device=rank)
+        
+        if cfg.warmup:
+            scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+                self.optimizer, cfg.nepoch - cfg.warmup_epoch, eta_min=1e-6
+            )
+            self.scheduler = GradualWarmupScheduler(
+                self.optimizer, multiplier=1, total_epoch=cfg.warmup_epoch, after_scheduler=scheduler_cosine
+            )
+            self.scheduler.step()        
+
+    def adjust_learning_rate(self, min_lr=1e-6):
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] /= 10.0
+            if param_group["lr"] < min_lr:
+                return False
+        return True
+
+
+    def set_input(self, input, istrain=True):
+        if self.reproduce_dire:
+            dire, label = input
+            self.dire = dire.to(self.device)
+            self.label = label.to(self.device).float()
+            
+        else:
+            img, dire, eps, label = input # if len(input) == 3 else (input[0], input[1], {})
+            H, W = img.shape[-2:]
+            B = img.shape[0]
+
+            # random jpeg compression
+            if istrain:
+                comp_quality = torch.randint(10, 100, (B,))
+                img = (img+1)/2
+                img = (img*255).to(torch.uint8)
+                img = torch.stack([decode_jpeg(encode_jpeg(img[i], quality=int(comp_quality[i]))) for i in range(B)], dim=0)
+                img = img / 255.
+                img = (img*2)-1
+            # only-img
+            if self.only_img:
+                img = img.to(self.device)
+                # calc eps from img
+                eps = dire_get_first_step_noise(img, self.adm, self.diffusion, self.adm_args, self.device)
+
+            # cutmix
+            # if torch.rand(1) < 0.3 and istrain:
+            #     c_lambda = torch.rand(1)
+            #     r_x = torch.randint(0, W, (1,))
+            #     r_y = torch.randint(0, H, (1,))
+            #     r_w = int(torch.sqrt(1-c_lambda)*W)
+            #     r_h = int(torch.sqrt(1-c_lambda)*H)
+
+            #     img[:, :, r_y:r_y+r_h, r_x:r_x+r_w] = img[0:1, :, r_y:r_y+r_h, r_x:r_x+r_w].repeat(B, 1, 1, 1)
+            #     dire[:, :, r_y:r_y+r_h, r_x:r_x+r_w] = dire[0:1, :, r_y:r_y+r_h, r_x:r_x+r_w].repeat(B, 1, 1, 1)
+            #     eps[:, :, r_y:r_y+r_h, r_x:r_x+r_w] = eps[0:1, :, r_y:r_y+r_h, r_x:r_x+r_w].repeat(B, 1, 1, 1)
+            #     label = c_lambda * label + (1-c_lambda) * label[0:1]
+            self.input = img.to(self.device)
+            self.dire = dire.to(self.device)
+            self.eps = eps.to(self.device)
+            self.label = label.to(self.device).float()
+       
+   
+    def forward(self):
+        if self.reproduce_dire:
+            self.output = self.student(self.dire)
+        else:
+            if self.only_eps or self.only_img:
+                self.output = self.student(self.eps)
+            else:
+                self.output = self.student(self.input, self.eps)
+            with torch.no_grad():
+                self.teacher_feature = self.teacher(self.dire)
+
+
+    def get_loss(self, kd=True):
+        loss = self.cls_criterion(self.output['logit'].squeeze(), self.label) 
+        if kd and (not self.reproduce_dire):
+            loss2 = self.kd_criterion(self.output['feature'], self.teacher_feature)
+            loss = loss + loss2 * self.kd_weight
+        return loss
+
+
+    def optimize_parameters(self):
+        self.optimizer.zero_grad()
+        self.forward()
+        self.loss = self.get_loss(self.kd)
+        self.loss.backward()
+        self.optimizer.step()
+
+    
+    
+    def load_networks(self, model_path):
+        state_dict = torch.load(model_path, map_location=self.device)
+        model_state_dict = state_dict["model"]
+        optimizer_state_dict = state_dict["optimizer"]
+        model_state_dict = {k.replace("module.", ""): v for k, v in model_state_dict.items()}
+        optimizer_state_dict = {k.replace("module.", ""): v for k, v in optimizer_state_dict.items()}
+
+        if self.distributed:
+            self.student.module.load_state_dict(model_state_dict)
+        else:
+            self.student.load_state_dict(model_state_dict)
+        self.optimizer.load_state_dict(optimizer_state_dict)
+        
+        print(f"Model loaded from {model_path}")
+        return True 
+    
+    @torch.no_grad()
+    def validate(self, gather=False, save=False, save_name=""):
+        self.student.eval()
+        y_pred = []
+        y_true = []
+        N_FAKE, N_REAL = 0, 0
+        for data, path in tqdm(self.val_loader, desc=f"Validation after {self.cur_epoch} epoch..."):
+            self.set_input(data, istrain=False)
+            self.forward()
+            pred = self.output['logit'].sigmoid()
+            
+            if gather:
+                try:
+                    dist
+                except:
+                    import torch.distributed as dist
+                pred_gather = [pred for _ in range(self.world_size)]
+                label_gather = [self.label for _ in range(self.world_size)]
+                dist.all_gather(pred_gather, pred)
+                dist.all_gather(label_gather, self.label)
+            else:
+                pred_gather = [pred]
+                label_gather = [self.label]
+            N_FAKE += sum([(label == 1).sum().item() for label in label_gather])
+            N_REAL += sum([(label == 0).sum().item() for label in label_gather])
+            
+            y_pred.extend(torch.cat(pred_gather).flatten().detach().cpu().tolist())
+            y_true.extend(torch.cat(label_gather).flatten().detach().cpu().tolist())
+        
+        y_true, y_pred = np.array(y_true), np.array(y_pred)
+        acc = accuracy_score(y_true, y_pred > 0.5)
+        ap = average_precision_score(y_true, y_pred)
+        precision = precision_score(y_true, y_pred > 0.5)
+        if self.run:
+            self.run.log({"val_acc": acc, "val_ap": ap , "val_precision": precision})
+            self.run.log({"N_FAKE": N_FAKE, "N_REAL": N_REAL})
+        print(f"Validation: acc: {acc}, ap: {ap}, precision: {precision}")
+        print(f"N_FAKE: {N_FAKE}, N_REAL: {N_REAL}")
+        if save:
+            with open(save_name, "w") as f:
+                f.write(f"Validation: acc: {acc}, ap: {ap}, precision: {precision}\n")
+                f.write(f"N_FAKE: {N_FAKE}, N_REAL: {N_REAL}")
+        
+        
+    def train(self):
+        for epoch in range(self.cfg.nepoch):
+            if self.run:
+                self.run.log({"epoch": epoch})
+            if epoch % self.val_every == 0 and epoch != 0:
+                self.validate(gather=self.distributed)
+            
+            self.student.train()
+            for data_and_paths in tqdm(self.train_loader, desc=f"Trainig {epoch} epoch..."):
+                data, path = data_and_paths
+                self.total_steps += 1
+                self.set_input(data)
+                self.optimize_parameters()
+                if self.total_steps % 100 == 0 and self.run:
+                    print(f"step: {self.total_steps}, loss: {self.loss}")
+                    self.run.log({"loss": self.loss, "step": self.total_steps})
+            
+            if self.run:
+                self.save_networks(epoch)
+                
+            if self.cfg.warmup:
+                self.scheduler.step()
+            if self.distributed:
+                dist.barrier()
+            self.cur_epoch = epoch

utils/utils.py

@@ -0,0 +1,47 @@
+import argparse
+import warnings
+import numpy as np
+import torch
+
+
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.nn.functional")
+
+
+def str2bool(v: str, strict=True) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif isinstance(v, str):
+        if v.lower() in ("true", "yes", "on" "t", "y", "1"):
+            return True
+        elif v.lower() in ("false", "no", "off", "f", "n", "0"):
+            return False
+    if strict:
+        raise argparse.ArgumentTypeError("Unsupported value encountered.")
+    else:
+        return True
+
+
+def to_cuda(data, device="cuda", exclude_keys: "list[str]" = None):
+    if isinstance(data, torch.Tensor):
+        data = data.to(device)
+    elif isinstance(data, (tuple, list, set)):
+        data = [to_cuda(b, device) for b in data]
+    elif isinstance(data, dict):
+        if exclude_keys is None:
+            exclude_keys = []
+        for k in data.keys():
+            if k not in exclude_keys:
+                data[k] = to_cuda(data[k], device)
+    else:
+        # raise TypeError(f"Unsupported type: {type(data)}")
+        data = data
+    return data
+
+
+def pad_img_to_square(img: np.ndarray):
+    H, W = img.shape[:2]
+    if H != W:
+        new_size = max(H, W)
+        img = np.pad(img, ((0, new_size - H), (0, new_size - W), (0, 0)), mode="constant")
+        assert img.shape[0] == img.shape[1] == new_size
+    return img

utils/warmup.py

@@ -0,0 +1,70 @@
+from torch.optim.lr_scheduler import ReduceLROnPlateau, _LRScheduler
+
+
+class GradualWarmupScheduler(_LRScheduler):
+    """Gradually warm-up(increasing) learning rate in optimizer.
+    Proposed in 'Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour'.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        multiplier: target learning rate = base lr * multiplier if multiplier > 1.0. if multiplier = 1.0, lr starts from 0 and ends up with the base_lr.
+        total_epoch: target learning rate is reached at total_epoch, gradually
+        after_scheduler: after target_epoch, use this scheduler(eg. ReduceLROnPlateau)
+    """
+
+    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):
+        self.multiplier = multiplier
+        if self.multiplier < 1.0:
+            raise ValueError("multiplier should be greater thant or equal to 1.")
+        self.total_epoch = total_epoch
+        self.after_scheduler = after_scheduler
+        self.finished = False
+        super().__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch > self.total_epoch:
+            if self.after_scheduler:
+                if not self.finished:
+                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]
+                    self.finished = True
+                return self.after_scheduler.get_last_lr()
+            return [base_lr * self.multiplier for base_lr in self.base_lrs]
+
+        if self.multiplier == 1.0:
+            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]
+        else:
+            return [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+
+    def step_ReduceLROnPlateau(self, metrics, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+        self.last_epoch = (
+            epoch if epoch != 0 else 1
+        )  # ReduceLROnPlateau is called at the end of epoch, whereas others are called at beginning
+        if self.last_epoch <= self.total_epoch:
+            warmup_lr = [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+            for param_group, lr in zip(self.optimizer.param_groups, warmup_lr):
+                param_group["lr"] = lr
+        else:
+            if epoch is None:
+                self.after_scheduler.step(metrics, None)
+            else:
+                self.after_scheduler.step(metrics, epoch - self.total_epoch)
+
+    def step(self, epoch=None, metrics=None):
+        if type(self.after_scheduler) != ReduceLROnPlateau:
+            if self.finished and self.after_scheduler:
+                if epoch is None:
+                    self.after_scheduler.step(None)
+                else:
+                    self.after_scheduler.step(epoch - self.total_epoch)
+            else:
+                return super().step(epoch)
+        else:
+            self.step_ReduceLROnPlateau(metrics, epoch)