[Update] Add files and checkpoint

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ttf/KaiXinSongA.ttf filter=lfs diff=lfs merge=lfs -text
+ttf/KaiXinSongB.ttf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,180 @@
+# Initially taken from GitHub's Python gitignore file
+outputs/
+run_sh/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# tests and logs
+tests/fixtures/cached_*_text.txt
+logs/
+lightning_logs/
+lang_code_data/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a Python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vs
+.vscode
+
+# Pycharm
+.idea
+
+# TF code
+tensorflow_code
+
+# Models
+proc_data
+
+# examples
+runs
+/runs_old
+/wandb
+/examples/runs
+/examples/**/*.args
+/examples/rag/sweep
+
+# data
+/data
+serialization_dir
+
+# emacs
+*.*~
+debug.env
+
+# vim
+.*.swp
+
+# ctags
+tags
+
+# pre-commit
+.pre-commit*
+
+# .lock
+*.lock
+
+# DS_Store (MacOS)
+.DS_Store
+
+# RL pipelines may produce mp4 outputs
+*.mp4
+
+# dependencies
+/transformers
+
+# ruff
+.ruff_cache
+
+# wandb
+wandb

app.py

@@ -0,0 +1,147 @@
+import random
+import gradio as gr
+from sample import (arg_parse, 
+                    sampling,
+                    load_fontdiffuer_pipeline)
+
+
+def run_fontdiffuer(source_image, 
+                    character, 
+                    reference_image,
+                    sampling_step,
+                    guidance_scale,
+                    batch_size):
+    args.character_input = False if source_image is not None else True
+    args.content_character = character
+    args.sampling_step = sampling_step
+    args.guidance_scale = guidance_scale
+    args.batch_size = batch_size
+    args.seed = random.randint(0, 10000)
+    out_image = sampling(
+        args=args,
+        pipe=pipe,
+        content_image=source_image,
+        style_image=reference_image)
+    return out_image
+
+
+if __name__ == '__main__':
+    args = arg_parse()
+    args.demo = True
+    args.ckpt_dir = 'ckpt'
+    args.ttf_path = 'ttf/KaiXinSongA.ttf'
+
+    # load fontdiffuer pipeline
+    pipe = load_fontdiffuer_pipeline(args=args)
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.HTML("""
+                    <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
+                    <h1 style="font-weight: 900; font-size: 3rem; margin: 0rem">
+                        FontDiffuser
+                    </h1>
+                    <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+                        <a href="https://yeungchenwa.github.io/"">Zhenhua Yang</a>, 
+                        <a href="https://scholar.google.com/citations?user=6zNgcjAAAAAJ&hl=zh-CN&oi=ao"">Dezhi Peng</a>, 
+                        Yuxin Kong, Yuyi Zhang, 
+                        <a href="https://scholar.google.com/citations?user=IpmnLFcAAAAJ&hl=zh-CN&oi=ao"">Cong Yao</a>, 
+                        <a href="http://www.dlvc-lab.net/lianwen/Index.html"">Lianwen Jin</a>†
+                    </h2>
+                    <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+                        <strong>South China University of Technology</strong>, Alibaba DAMO Academy
+                    </h2>
+                    <h3 style="font-weight: 450; font-size: 1rem; margin: 0rem"> 
+                    [<a href="https://github.com/yeungchenwa/FontDiffuser" style="color:blue;">arXiv</a>] 
+                    [<a href="https://github.com/yeungchenwa/FontDiffuser" style="color:green;">Github</a>]
+                    </h3>
+                    <h2 style="text-align: left; font-weight: 600; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
+                    1.We propose FontDiffuser, which is capable to generate unseen characters and styles, and it can be extended to the cross-lingual generation, such as Chinese to Korean.
+                    </h2>
+                    <h2 style="text-align: left; font-weight: 600; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
+                    2. FontDiffuser excels in generating complex character and handling large style variation. And it achieves state-of-the-art performance.
+                    </h2>
+                    </div>
+                    """)
+                gr.Image('figures/result_vis.png')
+                gr.Image('figures/demo_tips.png')
+            with gr.Column(scale=1):
+                with gr.Row():
+                    source_image = gr.Image(width=320, label='[Option 1] Source Image', image_mode='RGB', type='pil')
+                    reference_image = gr.Image(width=320, label='Reference Image', image_mode='RGB', type='pil')
+                with gr.Row():
+                    character = gr.Textbox(value='隆', label='[Option 2] Source Character')
+                with gr.Row():
+                    fontdiffuer_output_image = gr.Image(height=200, label="FontDiffuser Output Image", image_mode='RGB', type='pil')
+
+                sampling_step = gr.Slider(20, 50, value=20, step=10, 
+                                          label="Sampling Step", info="The sampling step by FontDiffuser.")
+                guidance_scale = gr.Slider(1, 12, value=7.5, step=0.5, 
+                                           label="Scale of Classifier-free Guidance", 
+                                           info="The scale used for classifier-free guidance sampling")
+                batch_size = gr.Slider(1, 4, value=1, step=1, 
+                                       label="Batch Size", info="The number of images to be sampled.")
+
+                FontDiffuser = gr.Button('Run FontDiffuser')
+                gr.Markdown("## <font color=#008000, size=6>Examples that You Can Choose Below⬇️</font>")
+        with gr.Row():
+            gr.Markdown("## Examples")
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("## Example 1️⃣: Source Image and Reference Image")
+                gr.Markdown("### In this mode, we provide both the source image and \
+                            the reference image for you to try our demo!")
+                gr.Examples(
+                    examples=[['figures/source_imgs/source_灨.jpg', 'figures/ref_imgs/ref_籍.jpg'], 
+                            ['figures/source_imgs/source_鑻.jpg', 'figures/ref_imgs/ref_鹰.jpg'],
+                            ['figures/source_imgs/source_鑫.jpg', 'figures/ref_imgs/ref_壤.jpg'],
+                            ['figures/source_imgs/source_釅.jpg', 'figures/ref_imgs/ref_雕.jpg']],
+                    inputs=[source_image, reference_image]
+                )
+            with gr.Column(scale=1):
+                gr.Markdown("## Example 2️⃣: Character and Reference Image")
+                gr.Markdown("### In this mode, we provide the content character and the reference image \
+                            for you to try our demo!")
+                gr.Examples(
+                    examples=[['龍', 'figures/ref_imgs/ref_鷢.jpg'],
+                            ['轉', 'figures/ref_imgs/ref_鲸.jpg'],
+                            ['懭', 'figures/ref_imgs/ref_籍_1.jpg'],
+                            ['識', 'figures/ref_imgs/ref_鞣.jpg']],
+                    inputs=[character, reference_image]
+                )
+            with gr.Column(scale=1):
+                gr.Markdown("## Example 3️⃣: Reference Image")
+                gr.Markdown("### In this mode, we provide only the reference image, \
+                            you can upload your own source image or you choose the character above \
+                            to try our demo!")
+                gr.Examples(
+                    examples=['figures/ref_imgs/ref_闡.jpg', 
+                            'figures/ref_imgs/ref_雕.jpg',
+                            'figures/ref_imgs/ref_豄.jpg',
+                            'figures/ref_imgs/ref_馨.jpg',
+                            'figures/ref_imgs/ref_鲸.jpg',
+                            'figures/ref_imgs/ref_檀.jpg',
+                            'figures/ref_imgs/ref_鞣.jpg',
+                            'figures/ref_imgs/ref_穗.jpg',
+                            'figures/ref_imgs/ref_欟.jpg',
+                            'figures/ref_imgs/ref_籍_1.jpg',
+                            'figures/ref_imgs/ref_鷢.jpg',
+                            'figures/ref_imgs/ref_媚.jpg',
+                            'figures/ref_imgs/ref_籍.jpg',
+                            'figures/ref_imgs/ref_壤.jpg',
+                            'figures/ref_imgs/ref_蜓.jpg',
+                            'figures/ref_imgs/ref_鹰.jpg'],
+                    examples_per_page=20,
+                    inputs=reference_image
+                )
+        FontDiffuser.click(
+            fn=run_fontdiffuer,
+            inputs=[source_image, 
+                    character, 
+                    reference_image,
+                    sampling_step,
+                    guidance_scale,
+                    batch_size],
+            outputs=fontdiffuer_output_image)
+    demo.launch(debug=True)

ckpt/content_encoder.pth

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

ckpt/style_encoder.pth

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

ckpt/unet.pth

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

configs/fontdiffuser.py

@@ -0,0 +1,87 @@
+import os
+import argparse
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Training config for FontDiffuser.")
+    ################# Experience #################
+    parser.add_argument("--seed", type=int, default=123, help="A seed for reproducible training.")
+    parser.add_argument("--experience_name", type=str, default="fontdiffuer_training")
+    parser.add_argument("--data_root", type=str, default=None, 
+                        help="The font dataset root path.",)
+    parser.add_argument("--output_dir", type=str, default=None, 
+                        help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--report_to", type=str, default="tensorboard")
+    parser.add_argument("--logging_dir", type=str, default="logs", 
+                        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+                              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."))
+
+    # Model
+    parser.add_argument("--resolution", type=int, default=96, 
+                        help="The resolution for input images, all the images in the train/validation \
+                            dataset will be resized to this.")
+    parser.add_argument("--unet_channels", type=tuple, default=(64, 128, 256, 512),
+                        help="The channels of the UNet.")
+    parser.add_argument("--style_image_size", type=int, default=96, help="The size of style images.")
+    parser.add_argument("--content_image_size", type=int, default=96, help="The size of content images.")
+    parser.add_argument("--content_encoder_downsample_size", type=int, default=3, 
+                        help="The downsample size of the content encoder.")
+    parser.add_argument("--channel_attn", type=bool, default=True, help="Whether to use the se attention.",)
+    parser.add_argument("--content_start_channel", type=int, default=64, 
+                        help="The channels of the fisrt layer output of content encoder.",)
+    parser.add_argument("--style_start_channel", type=int, default=64, 
+                        help="The channels of the fisrt layer output of content encoder.",)
+    
+    # Training
+    parser.add_argument("--train_batch_size", type=int, default=4, 
+                        help="Batch size (per device) for the training dataloader.")
+    ## loss coefficient
+    parser.add_argument("--perceptual_coefficient", type=float, default=0.01)
+    parser.add_argument("--offset_coefficient", type=float, default=0.5)
+    ## step
+    parser.add_argument("--max_train_steps", type=int, default=440000, 
+                        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",)
+    parser.add_argument("--ckpt_interval", type=int,default=40000, help="The step begin to validate.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, 
+                        help="Number of updates steps to accumulate before performing a backward/update pass.",)
+    parser.add_argument("--log_interval", type=int, default=100, help="The log interval of training.")
+    ## learning rate
+    parser.add_argument("--learning_rate", type=float, default=1e-4, 
+                        help="Initial learning rate (after the potential warmup period) to use.")
+    parser.add_argument("--scale_lr", action="store_true", default=False, 
+                        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
+    parser.add_argument("--lr_scheduler", type=str, default="linear", 
+                        help="The scheduler type to use. Choose between 'linear', 'cosine', \
+                            'cosine_with_restarts', 'polynomial', 'constant', 'constant_with_warmup'")
+    parser.add_argument("--lr_warmup_steps", type=int, default=10000, 
+                        help="Number of steps for the warmup in the lr scheduler.")
+    ## classifier-free
+    parser.add_argument("--drop_prob", type=float, default=0.1, help="The uncondition training drop out probability.")
+    ## scheduler
+    parser.add_argument("--beta_scheduler", type=str, default="scaled_linear", help="The beta scheduler for DDPM.")
+    ## optimizer
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    
+    parser.add_argument("--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"], 
+                        help="Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires \
+                            PyTorch >= 1.10. and an Nvidia Ampere GPU.")
+    
+    # Sampling
+    parser.add_argument("--algorithm_type", type=str, default="dpmsolver++", help="Algorithm for sampleing.")
+    parser.add_argument("--guidance_type", type=str, default="classifier-free", help="Guidance type of sampling.")
+    parser.add_argument("--guidance_scale", type=float, default=7.5, help="Guidance scale of the classifier-free mode.")
+    parser.add_argument("--num_inference_steps", type=int, default=20, help="Sampling step.")
+    parser.add_argument("--model_type", type=str, default="noise", help="model_type for sampling.")
+    parser.add_argument("--order", type=int, default=2, help="The order of the dpmsolver.")
+    parser.add_argument("--skip_type", type=str, default="time_uniform", help="Skip type of dpmsolver.")
+    parser.add_argument("--method", type=str, default="multistep", help="Multistep of dpmsolver.")
+    parser.add_argument("--correcting_x0_fn", type=str, default=None, help="correcting_x0_fn of dpmsolver.")
+    parser.add_argument("--t_start", type=str, default=None, help="t_start of dpmsolver.")
+    parser.add_argument("--t_end", type=str, default=None, help="t_end of dpmsolver.")
+    
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    
+    return parser

dataset/font_dataset.py

@@ -0,0 +1,69 @@
+import os
+import random
+from PIL import Image
+
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+
+def get_nonorm_transform(resolution):
+    nonorm_transform =  transforms.Compose(
+            [transforms.Resize((resolution, resolution), 
+                               interpolation=transforms.InterpolationMode.BILINEAR), 
+             transforms.ToTensor()])
+    return nonorm_transform
+
+
+class FontDataset(Dataset):
+    """The dataset of font generation  
+    """
+    def __init__(self, args, phase, transforms=None):
+        super().__init__()
+        self.root = args.data_root
+        self.phase = phase
+        
+        # Get Data path
+        self.get_path()
+        self.transforms = transforms
+        self.nonorm_transforms = get_nonorm_transform(args.resolution)
+
+    def get_path(self):
+        self.target_images = []
+        # images with related style  
+        self.style_to_images = {}
+        target_image_dir = f"{self.root}/{self.phase}/TargetImage"
+        for style in os.listdir(target_image_dir):
+            images_related_style = []
+            for img in os.listdir(f"{target_image_dir}/{style}"):
+                img_path = f"{target_image_dir}/{style}/{img}"
+                self.target_images.append(img_path)
+                images_related_style.append(img_path)
+            self.style_to_images[style] = images_related_style
+
+    def __getitem__(self, index):
+        target_image_path = self.target_images[index]
+        target_image_name = target_image_path.split('/')[-1]
+        style, content = target_image_name.split('.')[0].split('+')
+        
+        # Read content image
+        content_image_path = f"{self.root}/{self.phase}/ContentImage/{content}.jpg"
+        content_image = Image.open(content_image_path).convert('RGB')
+
+        # Random sample used for style image
+        images_related_style = self.style_to_images[style].copy()
+        images_related_style.remove(target_image_path)
+        style_image_path = random.choice(images_related_style)
+        style_image = Image.open(style_image_path).convert("RGB")
+        
+        # Read target image
+        target_image = Image.open(target_image_path).convert("RGB")
+        nonorm_target_image = self.nonorm_transforms(target_image)
+
+        if self.transforms is not None:
+            content_image = self.transforms[0](content_image)
+            style_image = self.transforms[1](style_image)
+            target_image = self.transforms[2](target_image)
+        
+        return content_image, style_image, target_image, nonorm_target_image, target_image_path
+
+    def __len__(self):
+        return len(self.target_images)

requirements.txt

@@ -0,0 +1,5 @@
+transformers==4.33.1
+accelerate==0.23.0
+diffusers==0.22.0.dev0
+gradio==4.8.0
+yaml

sample.py

@@ -0,0 +1,252 @@
+import os
+import cv2
+import time
+import random
+import numpy as np
+from PIL import Image
+
+import torch
+import torchvision.transforms as transforms
+from accelerate.utils import set_seed
+
+from src import (FontDiffuserDPMPipeline,
+                 FontDiffuserModelDPM,
+                 build_ddpm_scheduler,
+                 build_unet,
+                 build_content_encoder,
+                 build_style_encoder)
+from utils import (ttf2im,
+                   load_ttf,
+                   is_char_in_font,
+                   save_args_to_yaml,
+                   save_single_image,
+                   save_image_with_content_style)
+
+
+def arg_parse():
+    from configs.fontdiffuser import get_parser
+
+    parser = get_parser()
+    parser.add_argument("--ckpt_dir", type=str, default=None)
+    parser.add_argument("--demo", action="store_true")
+    parser.add_argument("--controlnet", type=bool, default=False, 
+                        help="If in demo mode, the controlnet can be added.")
+    parser.add_argument("--character_input", action="store_true")
+    parser.add_argument("--content_character", type=str, default=None)
+    parser.add_argument("--content_image_path", type=str, default=None)
+    parser.add_argument("--style_image_path", type=str, default=None)
+    parser.add_argument("--save_image", action="store_true")
+    parser.add_argument("--save_image_dir", type=str, default=None,
+                        help="The saving directory.")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--ttf_path", type=str, default="ttf/KaiXinSongA.ttf")
+    args = parser.parse_args()
+    style_image_size = args.style_image_size
+    content_image_size = args.content_image_size
+    args.style_image_size = (style_image_size, style_image_size)
+    args.content_image_size = (content_image_size, content_image_size)
+
+    return args
+
+
+def image_process(args, content_image=None, style_image=None):
+    if not args.demo:
+        # Read content image and style image
+        if args.character_input:
+            assert args.content_character is not None, "The content_character should not be None."
+            if not is_char_in_font(font_path=args.ttf_path, char=args.content_character):
+                return None, None
+            font = load_ttf(ttf_path=args.ttf_path)
+            content_image = ttf2im(font=font, char=args.content_character)
+            content_image_pil = content_image.copy()
+        else:
+            content_image = Image.open(args.content_image_path).convert('RGB')
+            content_image_pil = None
+        style_image = Image.open(args.style_image_path).convert('RGB')
+    else:
+        assert style_image is not None, "The style image should not be None."
+        if args.character_input:
+            assert args.content_character is not None, "The content_character should not be None."
+            if not is_char_in_font(font_path=args.ttf_path, char=args.content_character):
+                return None, None
+            font = load_ttf(ttf_path=args.ttf_path)
+            content_image = ttf2im(font=font, char=args.content_character)
+        else:
+            assert content_image is not None, "The content image should not be None."
+        content_image_pil = None
+        
+    ## Dataset transform
+    content_inference_transforms = transforms.Compose(
+        [transforms.Resize(args.content_image_size, \
+                            interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])])
+    style_inference_transforms = transforms.Compose(
+        [transforms.Resize(args.style_image_size, \
+                           interpolation=transforms.InterpolationMode.BILINEAR),
+         transforms.ToTensor(),
+         transforms.Normalize([0.5], [0.5])])
+    content_image = content_inference_transforms(content_image)[None, :]
+    style_image = style_inference_transforms(style_image)[None, :]
+
+    return content_image, style_image, content_image_pil
+
+def load_fontdiffuer_pipeline(args):
+    # Load the model state_dict
+    unet = build_unet(args=args)
+    unet.load_state_dict(torch.load(f"{args.ckpt_dir}/unet.pth"))
+    style_encoder = build_style_encoder(args=args)
+    style_encoder.load_state_dict(torch.load(f"{args.ckpt_dir}/style_encoder.pth"))
+    content_encoder = build_content_encoder(args=args)
+    content_encoder.load_state_dict(torch.load(f"{args.ckpt_dir}/content_encoder.pth"))
+    model = FontDiffuserModelDPM(
+        unet=unet,
+        style_encoder=style_encoder,
+        content_encoder=content_encoder)
+    model.to(args.device)
+    print("Loaded the model state_dict successfully!")
+
+    # Load the training ddpm_scheduler.
+    train_scheduler = build_ddpm_scheduler(args=args)
+    print("Loaded training DDPM scheduler sucessfully!")
+
+    # Load the DPM_Solver to generate the sample.
+    pipe = FontDiffuserDPMPipeline(
+        model=model,
+        ddpm_train_scheduler=train_scheduler,
+        model_type=args.model_type,
+        guidance_type=args.guidance_type,
+        guidance_scale=args.guidance_scale,
+    )
+    print("Loaded dpm_solver pipeline sucessfully!")
+
+    return pipe
+
+
+def sampling(args, pipe, content_image=None, style_image=None):
+    if not args.demo:
+        os.makedirs(args.save_image_dir, exist_ok=True)
+        # saving sampling config
+        save_args_to_yaml(args=args, output_file=f"{args.save_image_dir}/sampling_config.yaml")
+
+    if args.seed:
+        set_seed(seed=args.seed)
+    
+    content_image, style_image, content_image_pil = image_process(args=args, 
+                                                                  content_image=content_image, 
+                                                                  style_image=style_image)
+    if content_image == None:
+        print(f"The content_character you provided is not in the ttf. \
+                Please change the content_character or you can change the ttf.")
+        return None
+
+    with torch.no_grad():
+        content_image = content_image.to(args.device)
+        style_image = style_image.to(args.device)
+        print(f"Sampling by DPM-Solver++ ......")
+        start = time.time()
+        images = pipe.generate(
+            content_images=content_image,
+            style_images=style_image,
+            batch_size=1,
+            order=args.order,
+            num_inference_step=args.num_inference_steps,
+            content_encoder_downsample_size=args.content_encoder_downsample_size,
+            t_start=args.t_start,
+            t_end=args.t_end,
+            dm_size=args.content_image_size,
+            algorithm_type=args.algorithm_type,
+            skip_type=args.skip_type,
+            method=args.method,
+            correcting_x0_fn=args.correcting_x0_fn)
+        end = time.time()
+
+        if args.save_image:
+            print(f"Saving the image ......")
+            save_single_image(save_dir=args.save_image_dir, image=images[0])
+            if args.character_input:
+                save_image_with_content_style(save_dir=args.save_image_dir,
+                                            image=images[0],
+                                            content_image_pil=content_image_pil,
+                                            content_image_path=None,
+                                            style_image_path=args.style_image_path,
+                                            resolution=args.resolution)
+            else:
+                save_image_with_content_style(save_dir=args.save_image_dir,
+                                            image=images[0],
+                                            content_image_pil=None,
+                                            content_image_path=args.content_image_path,
+                                            style_image_path=args.style_image_path,
+                                            resolution=args.resolution)
+            print(f"Finish the sampling process, costing time {end - start}s")
+        return images[0]
+
+
+def load_controlnet_pipeline(args,
+                             config_path="lllyasviel/sd-controlnet-canny", 
+                             ckpt_path="runwayml/stable-diffusion-v1-5"):
+    from diffusers import ControlNetModel, AutoencoderKL
+    # load controlnet model and pipeline
+    from diffusers import StableDiffusionControlNetPipeline, UniPCMultistepScheduler
+    controlnet = ControlNetModel.from_pretrained(config_path, 
+                                                 torch_dtype=torch.float16,
+                                                 cache_dir=f"{args.ckpt_dir}/controlnet")
+    print(f"Loaded ControlNet Model Successfully!")
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(ckpt_path, 
+                                                             controlnet=controlnet, 
+                                                             torch_dtype=torch.float16,
+                                                             cache_dir=f"{args.ckpt_dir}/controlnet_pipeline")
+    # faster
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_model_cpu_offload()
+    print(f"Loaded ControlNet Pipeline Successfully!")
+
+    return pipe
+
+
+def controlnet(text_prompt, 
+               pil_image,
+               pipe):
+    image = np.array(pil_image)
+    # get canny image
+    image = cv2.Canny(image=image, threshold1=100, threshold2=200)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    canny_image = Image.fromarray(image)
+    
+    seed = random.randint(0, 10000)
+    generator = torch.manual_seed(seed)
+    image = pipe(text_prompt, 
+                 num_inference_steps=50, 
+                 generator=generator, 
+                 image=canny_image,
+                 output_type='pil').images[0]
+    return image
+
+
+def load_instructpix2pix_pipeline(args,
+                                  ckpt_path="timbrooks/instruct-pix2pix"):
+    from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
+    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(ckpt_path, 
+                                                                  torch_dtype=torch.float16)
+    pipe.to(args.device)
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+    return pipe
+
+def instructpix2pix(pil_image, text_prompt, pipe):
+    image = pil_image.resize((512, 512))
+    seed = random.randint(0, 10000)
+    generator = torch.manual_seed(seed)
+    image = pipe(prompt=text_prompt, image=image, generator=generator, 
+                 num_inference_steps=20, image_guidance_scale=1.1).images[0]
+
+    return image
+
+
+if __name__=="__main__":
+    args = arg_parse()
+    
+    # load fontdiffuser pipeline
+    pipe = load_fontdiffuer_pipeline(args=args)
+    out_image = sampling(args=args, pipe=pipe)

src/__init__.py

@@ -0,0 +1,11 @@
+from .model import (FontDiffuserModel,
+                   FontDiffuserModelDPM)
+from .criterion import ContentPerceptualLoss
+from .dpm_solver.pipeline_dpm_solver import FontDiffuserDPMPipeline
+from .modules import (ContentEncoder,
+                     StyleEncoder, 
+                     UNet)
+from .build import (build_unet, 
+                   build_ddpm_scheduler, 
+                   build_style_encoder, 
+                   build_content_encoder)

src/build.py

@@ -0,0 +1,64 @@
+from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
+from src import (ContentEncoder, 
+                 StyleEncoder, 
+                 UNet)
+
+
+def build_unet(args):
+    unet = UNet(
+        sample_size=args.resolution,
+        in_channels=3,
+        out_channels=3,
+        flip_sin_to_cos=True,
+        freq_shift=0,
+        down_block_types=('DownBlock2D', 
+                          'MCADownBlock2D',
+                          'MCADownBlock2D', 
+                          'DownBlock2D'),
+        up_block_types=('UpBlock2D', 
+                        'StyleRSIUpBlock2D',
+                        'StyleRSIUpBlock2D', 
+                        'UpBlock2D'),
+        block_out_channels=args.unet_channels, 
+        layers_per_block=2,
+        downsample_padding=1,
+        mid_block_scale_factor=1,
+        act_fn='silu',
+        norm_num_groups=32,
+        norm_eps=1e-05,
+        cross_attention_dim=args.style_start_channel * 16,
+        attention_head_dim=1,
+        channel_attn=args.channel_attn,
+        content_encoder_downsample_size=args.content_encoder_downsample_size,
+        content_start_channel=args.content_start_channel,
+        reduction=32)
+    
+    return unet
+
+
+def build_style_encoder(args):
+    style_image_encoder = StyleEncoder(
+        G_ch=args.style_start_channel,
+        resolution=args.style_image_size[0])
+    print("Get CG-GAN Style Encoder!")
+    return style_image_encoder
+
+
+def build_content_encoder(args):
+    content_image_encoder = ContentEncoder(
+        G_ch=args.content_start_channel,
+        resolution=args.content_image_size[0])
+    print("Get CG-GAN Content Encoder!")
+    return content_image_encoder
+
+
+def build_ddpm_scheduler(args):
+    ddpm_scheduler = DDPMScheduler(
+        num_train_timesteps=1000,
+        beta_start=0.0001,
+        beta_end=0.02,
+        beta_schedule=args.beta_scheduler,
+        trained_betas=None,
+        variance_type="fixed_small",
+        clip_sample=True)
+    return ddpm_scheduler

src/criterion.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torchvision 
+
+
+class VGG16(nn.Module):
+    def __init__(self):
+        super(VGG16, self).__init__()
+        vgg16 = torchvision.models.vgg16(pretrained=True)
+
+        self.enc_1 = nn.Sequential(*vgg16.features[:5])
+        self.enc_2 = nn.Sequential(*vgg16.features[5:10])
+        self.enc_3 = nn.Sequential(*vgg16.features[10:17])
+
+        for i in range(3):
+            for param in getattr(self, f'enc_{i+1:d}').parameters():
+                param.requires_grad = False
+
+    def forward(self, image):
+        results = [image]
+        for i in range(3):
+            func = getattr(self, f'enc_{i+1:d}')
+            results.append(func(results[-1]))
+        return results[1:]
+
+
+class ContentPerceptualLoss(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.VGG = VGG16()
+
+    def calculate_loss(self, generated_images, target_images, device):
+        self.VGG = self.VGG.to(device)
+
+        generated_features = self.VGG(generated_images)
+        target_features = self.VGG(target_images)
+
+        perceptual_loss = 0
+        perceptual_loss += torch.mean((target_features[0] - generated_features[0]) ** 2)
+        perceptual_loss += torch.mean((target_features[1] - generated_features[1]) ** 2)
+        perceptual_loss += torch.mean((target_features[2] - generated_features[2]) ** 2)
+        perceptual_loss /= 3
+        return perceptual_loss

src/dpm_solver/dpm_solver_pytorch.py

@@ -0,0 +1,1332 @@
+import torch
+import torch.nn.functional as F
+import math
+
+
+class NoiseScheduleVP:
+    def __init__(
+            self,
+            schedule='discrete',
+            betas=None,
+            alphas_cumprod=None,
+            continuous_beta_0=0.1,
+            continuous_beta_1=20.,
+            dtype=torch.float32,
+        ):
+        """Create a wrapper class for the forward SDE (VP type).
+
+        ***
+        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
+                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
+        ***
+
+        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
+        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
+        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
+
+            log_alpha_t = self.marginal_log_mean_coeff(t)
+            sigma_t = self.marginal_std(t)
+            lambda_t = self.marginal_lambda(t)
+
+        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
+
+            t = self.inverse_lambda(lambda_t)
+
+        ===============================================================
+
+        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
+
+        1. For discrete-time DPMs:
+
+            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
+                t_i = (i + 1) / N
+            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
+            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
+
+            Args:
+                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
+                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
+
+            Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+
+            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
+                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
+                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
+                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
+                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
+                and
+                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
+
+
+        2. For continuous-time DPMs:
+
+            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
+            schedule are the default settings in DDPM and improved-DDPM:
+
+            Args:
+                beta_min: A `float` number. The smallest beta for the linear schedule.
+                beta_max: A `float` number. The largest beta for the linear schedule.
+                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
+                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
+                T: A `float` number. The ending time of the forward process.
+
+        ===============================================================
+
+        Args:
+            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
+                    'linear' or 'cosine' for continuous-time DPMs.
+        Returns:
+            A wrapper object of the forward SDE (VP type).
+        
+        ===============================================================
+
+        Example:
+
+        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', betas=betas)
+
+        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
+
+        # For continuous-time DPMs (VPSDE), linear schedule:
+        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
+
+        """
+
+        if schedule not in ['discrete', 'linear', 'cosine']:
+            raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))
+
+        self.schedule = schedule
+        if schedule == 'discrete':
+            if betas is not None:
+                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
+            else:
+                assert alphas_cumprod is not None
+                log_alphas = 0.5 * torch.log(alphas_cumprod)
+            self.total_N = len(log_alphas)
+            self.T = 1.
+            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)).to(dtype=dtype)
+            self.log_alpha_array = log_alphas.reshape((1, -1,)).to(dtype=dtype)
+        else:
+            self.total_N = 1000
+            self.beta_0 = continuous_beta_0
+            self.beta_1 = continuous_beta_1
+            self.cosine_s = 0.008
+            self.cosine_beta_max = 999.
+            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
+            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
+            self.schedule = schedule
+            if schedule == 'cosine':
+                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
+                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
+                self.T = 0.9946
+            else:
+                self.T = 1.
+
+    def marginal_log_mean_coeff(self, t):
+        """
+        Compute log(alpha_t) of a given continuous-time label t in [0, T].
+        """
+        if self.schedule == 'discrete':
+            return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)).reshape((-1))
+        elif self.schedule == 'linear':
+            return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        elif self.schedule == 'cosine':
+            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
+            log_alpha_t =  log_alpha_fn(t) - self.cosine_log_alpha_0
+            return log_alpha_t
+
+    def marginal_alpha(self, t):
+        """
+        Compute alpha_t of a given continuous-time label t in [0, T].
+        """
+        return torch.exp(self.marginal_log_mean_coeff(t))
+
+    def marginal_std(self, t):
+        """
+        Compute sigma_t of a given continuous-time label t in [0, T].
+        """
+        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
+
+    def marginal_lambda(self, t):
+        """
+        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
+        """
+        log_mean_coeff = self.marginal_log_mean_coeff(t)
+        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
+        return log_mean_coeff - log_std
+
+    def inverse_lambda(self, lamb):
+        """
+        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
+        """
+        if self.schedule == 'linear':
+            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            Delta = self.beta_0**2 + tmp
+            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
+        elif self.schedule == 'discrete':
+            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
+            t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
+            return t.reshape((-1,))
+        else:
+            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
+            t = t_fn(log_alpha)
+            return t
+
+
+def model_wrapper(
+    model,
+    noise_schedule,
+    model_type="noise",
+    model_kwargs={},
+    guidance_type="uncond",
+    condition=None,
+    unconditional_condition=None,
+    guidance_scale=1.,
+    classifier_fn=None,
+    classifier_kwargs={},
+):
+    """Create a wrapper function for the noise prediction model.
+
+    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
+    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
+
+    We support four types of the diffusion model by setting `model_type`:
+
+        1. "noise": noise prediction model. (Trained by predicting noise).
+
+        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
+
+        3. "v": velocity prediction model. (Trained by predicting the velocity).
+            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
+
+            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
+                arXiv preprint arXiv:2202.00512 (2022).
+            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
+                arXiv preprint arXiv:2210.02303 (2022).
+    
+        4. "score": marginal score function. (Trained by denoising score matching).
+            Note that the score function and the noise prediction model follows a simple relationship:
+            ```
+                noise(x_t, t) = -sigma_t * score(x_t, t)
+            ```
+
+    We support three types of guided sampling by DPMs by setting `guidance_type`:
+        1. "uncond": unconditional sampling by DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            `` 
+
+            The input `classifier_fn` has the following format:
+            ``
+                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
+            ``
+
+            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
+                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
+
+        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
+            `` 
+            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
+
+            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
+                arXiv preprint arXiv:2207.12598 (2022).
+        
+
+    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
+    or continuous-time labels (i.e. epsilon to T).
+
+    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
+    ``
+        def model_fn(x, t_continuous) -> noise:
+            t_input = get_model_input_time(t_continuous)
+            return noise_pred(model, x, t_input, **model_kwargs)         
+    ``
+    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
+
+    ===============================================================
+
+    Args:
+        model: A diffusion model with the corresponding format described above.
+        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+        model_type: A `str`. The parameterization type of the diffusion model.
+                    "noise" or "x_start" or "v" or "score".
+        model_kwargs: A `dict`. A dict for the other inputs of the model function.
+        guidance_type: A `str`. The type of the guidance for sampling.
+                    "uncond" or "classifier" or "classifier-free".
+        condition: A pytorch tensor. The condition for the guided sampling.
+                    Only used for "classifier" or "classifier-free" guidance type.
+        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
+                    Only used for "classifier-free" guidance type.
+        guidance_scale: A `float`. The scale for the guided sampling.
+        classifier_fn: A classifier function. Only used for the classifier guidance.
+        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
+    Returns:
+        A noise prediction model that accepts the noised data and the continuous time as the inputs.
+    """
+
+    def get_model_input_time(t_continuous):
+        """
+        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
+        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
+        For continuous-time DPMs, we just use `t_continuous`.
+        """
+        if noise_schedule.schedule == 'discrete':
+            return (t_continuous - 1. / noise_schedule.total_N) * 1000.
+        else:
+            return t_continuous
+
+    def noise_pred_fn(x, t_continuous, cond=None):
+        t_input = get_model_input_time(t_continuous)
+        if cond is None:
+            output = model(x, t_input, **model_kwargs)
+        else:
+            output = model(x, t_input, cond, **model_kwargs)
+        if model_type == "noise":
+            return output
+        elif model_type == "x_start":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            return (x - alpha_t * output) / sigma_t
+        elif model_type == "v":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            return alpha_t * output + sigma_t * x
+        elif model_type == "score":
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            return -sigma_t * output
+
+    def cond_grad_fn(x, t_input):
+        """
+        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
+        """
+        with torch.enable_grad():
+            x_in = x.detach().requires_grad_(True)
+            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
+            return torch.autograd.grad(log_prob.sum(), x_in)[0]
+
+    def model_fn(x, t_continuous):
+        """
+        The noise predicition model function that is used for DPM-Solver.
+        """
+        if guidance_type == "uncond":
+            return noise_pred_fn(x, t_continuous)
+        elif guidance_type == "classifier":
+            assert classifier_fn is not None
+            t_input = get_model_input_time(t_continuous)
+            cond_grad = cond_grad_fn(x, t_input)
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            noise = noise_pred_fn(x, t_continuous)
+            return noise - guidance_scale * sigma_t * cond_grad
+        elif guidance_type == "classifier-free":
+            if guidance_scale == 1. or unconditional_condition is None:
+                return noise_pred_fn(x, t_continuous, cond=condition)
+            elif model_kwargs["version"] == "V1" or model_kwargs["version"] == "V2_ConStyle" or model_kwargs["version"] == "V3":  # add this
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                c_in = []
+                c_in.append(torch.cat([unconditional_condition[0], condition[0]], dim=0))
+                c_in.append(torch.cat([unconditional_condition[1], condition[1]], dim=0))
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+                return noise_uncond + guidance_scale * (noise - noise_uncond)
+            elif model_kwargs["version"] == "FG_Sep":
+                x_in = torch.cat([x] * 3)
+                t_in = torch.cat([t_continuous] * 3)
+                c_in = []
+                c_in.append(torch.cat([unconditional_condition[0], unconditional_condition[0], condition[0]], dim=0))
+                c_in.append(torch.cat([unconditional_condition[1], condition[1], unconditional_condition[1]], dim=0))
+                noise_uncond, noise_cond_style, noise_cond_content = noise_pred_fn(x_in, t_in, cond=c_in).chunk(3)
+
+                style_guidance_scale = guidance_scale[0]
+                content_guidance_scale = guidance_scale[1]
+                return noise_uncond + style_guidance_scale * (noise_cond_style - noise_uncond) + content_guidance_scale * (noise_cond_content - noise_uncond)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                c_in = torch.cat([unconditional_condition, condition])
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+                return noise_uncond + guidance_scale * (noise - noise_uncond)
+
+    assert model_type in ["noise", "x_start", "v"]
+    assert guidance_type in ["uncond", "classifier", "classifier-free"]
+    return model_fn
+
+
+class DPM_Solver:
+    def __init__(
+        self,
+        model_fn,
+        noise_schedule,
+        algorithm_type="dpmsolver++",
+        correcting_x0_fn=None,
+        correcting_xt_fn=None,
+        thresholding_max_val=1.,
+        dynamic_thresholding_ratio=0.995,
+    ):
+        """Construct a DPM-Solver. 
+
+        We support both DPM-Solver (`algorithm_type="dpmsolver"`) and DPM-Solver++ (`algorithm_type="dpmsolver++"`).
+
+        We also support the "dynamic thresholding" method in Imagen[1]. For pixel-space diffusion models, you
+        can set both `algorithm_type="dpmsolver++"` and `correcting_x0_fn="dynamic_thresholding"` to use the
+        dynamic thresholding. The "dynamic thresholding" can greatly improve the sample quality for pixel-space
+        DPMs with large guidance scales. Note that the thresholding method is **unsuitable** for latent-space
+        DPMs (such as stable-diffusion).
+
+        To support advanced algorithms in image-to-image applications, we also support corrector functions for
+        both x0 and xt.
+
+        Args:
+            model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
+                ``
+                def model_fn(x, t_continuous):
+                    return noise
+                ``
+                The shape of `x` is `(batch_size, **shape)`, and the shape of `t_continuous` is `(batch_size,)`.
+            noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+            algorithm_type: A `str`. Either "dpmsolver" or "dpmsolver++".
+            correcting_x0_fn: A `str` or a function with the following format:
+                ```
+                def correcting_x0_fn(x0, t):
+                    x0_new = ...
+                    return x0_new
+                ```
+                This function is to correct the outputs of the data prediction model at each sampling step. e.g.,
+                ```
+                x0_pred = data_pred_model(xt, t)
+                if correcting_x0_fn is not None:
+                    x0_pred = correcting_x0_fn(x0_pred, t)
+                xt_1 = update(x0_pred, xt, t)
+                ```
+                If `correcting_x0_fn="dynamic_thresholding"`, we use the dynamic thresholding proposed in Imagen[1].
+            correcting_xt_fn: A function with the following format:
+                ```
+                def correcting_xt_fn(xt, t, step):
+                    x_new = ...
+                    return x_new
+                ```
+                This function is to correct the intermediate samples xt at each sampling step. e.g.,
+                ```
+                xt = ...
+                xt = correcting_xt_fn(xt, t, step)
+                ```
+            thresholding_max_val: A `float`. The max value for thresholding.
+                Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`.
+            dynamic_thresholding_ratio: A `float`. The ratio for dynamic thresholding (see Imagen[1] for details).
+                Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`.
+
+        [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour,
+            Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models
+            with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
+        """
+        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
+        self.noise_schedule = noise_schedule
+        assert algorithm_type in ["dpmsolver", "dpmsolver++"]
+        self.algorithm_type = algorithm_type
+        if correcting_x0_fn == "dynamic_thresholding":
+            self.correcting_x0_fn = self.dynamic_thresholding_fn
+        else:
+            self.correcting_x0_fn = correcting_x0_fn
+        self.correcting_xt_fn = correcting_xt_fn
+        self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
+        self.thresholding_max_val = thresholding_max_val
+
+    def dynamic_thresholding_fn(self, x0):
+        """
+        The dynamic thresholding method. 
+        """
+        dims = x0.dim()
+        p = self.dynamic_thresholding_ratio
+        s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+        s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
+        x0 = torch.clamp(x0, -s, s) / s
+        return x0
+
+    def noise_prediction_fn(self, x, t):
+        """
+        Return the noise prediction model.
+        """
+        return self.model(x, t)
+
+    def data_prediction_fn(self, x, t):
+        """
+        Return the data prediction model (with corrector).
+        """
+        noise = self.noise_prediction_fn(x, t)
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
+        x0 = (x - sigma_t * noise) / alpha_t
+        if self.correcting_x0_fn is not None:
+            x0 = self.correcting_x0_fn(x0)
+        return x0
+
+    def model_fn(self, x, t):
+        """
+        Convert the model to the noise prediction model or the data prediction model. 
+        """
+        if self.algorithm_type == "dpmsolver++":
+            return self.data_prediction_fn(x, t)
+        else:
+            return self.noise_prediction_fn(x, t)
+
+    def get_time_steps(self, skip_type, t_T, t_0, N, device):
+        """Compute the intermediate time steps for sampling.
+
+        Args:
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            N: A `int`. The total number of the spacing of the time steps.
+            device: A torch device.
+        Returns:
+            A pytorch tensor of the time steps, with the shape (N + 1,).
+        """
+        if skip_type == 'logSNR':
+            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
+            return self.noise_schedule.inverse_lambda(logSNR_steps)
+        elif skip_type == 'time_uniform':
+            return torch.linspace(t_T, t_0, N + 1).to(device)
+        elif skip_type == 'time_quadratic':
+            t_order = 2
+            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
+            return t
+        else:
+            raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
+
+    def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
+        """
+        Get the order of each step for sampling by the singlestep DPM-Solver.
+
+        We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
+        Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
+            - If order == 1:
+                We take `steps` of DPM-Solver-1 (i.e. DDIM).
+            - If order == 2:
+                - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
+                - If steps % 2 == 0, we use K steps of DPM-Solver-2.
+                - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
+            - If order == 3:
+                - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
+
+        ============================================
+        Args:
+            order: A `int`. The max order for the solver (2 or 3).
+            steps: A `int`. The total number of function evaluations (NFE).
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            device: A torch device.
+        Returns:
+            orders: A list of the solver order of each step.
+        """
+        if order == 3:
+            K = steps // 3 + 1
+            if steps % 3 == 0:
+                orders = [3,] * (K - 2) + [2, 1]
+            elif steps % 3 == 1:
+                orders = [3,] * (K - 1) + [1]
+            else:
+                orders = [3,] * (K - 1) + [2]
+        elif order == 2:
+            if steps % 2 == 0:
+                K = steps // 2
+                orders = [2,] * K
+            else:
+                K = steps // 2 + 1
+                orders = [2,] * (K - 1) + [1]
+        elif order == 1:
+            K = 1
+            orders = [1,] * steps
+        else:
+            raise ValueError("'order' must be '1' or '2' or '3'.")
+        if skip_type == 'logSNR':
+            # To reproduce the results in DPM-Solver paper
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+        else:
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
+        return timesteps_outer, orders
+
+    def denoise_to_zero_fn(self, x, s):
+        """
+        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization. 
+        """
+        return self.data_prediction_fn(x, s)
+
+    def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
+        """
+        DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        dims = x.dim()
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = (
+                sigma_t / sigma_s * x
+                - alpha_t * phi_1 * model_s
+            )
+            if return_intermediate:
+                return x_t, {'model_s': model_s}
+            else:
+                return x_t
+        else:
+            phi_1 = torch.expm1(h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = (
+                torch.exp(log_alpha_t - log_alpha_s) * x
+                - (sigma_t * phi_1) * model_s
+            )
+            if return_intermediate:
+                return x_t, {'model_s': model_s}
+            else:
+                return x_t
+
+    def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpmsolver'):
+        """
+        Singlestep solver DPM-Solver-2 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            r1: A `float`. The hyperparameter of the second-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpmsolver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type))
+        if r1 is None:
+            r1 = 0.5
+        ns = self.noise_schedule
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
+        alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_11 = torch.expm1(-r1 * h)
+            phi_1 = torch.expm1(-h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (
+                (sigma_s1 / sigma_s) * x
+                - (alpha_s1 * phi_11) * model_s
+            )
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    - (0.5 / r1) * (alpha_t * phi_1) * (model_s1 - model_s)
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1. / r1) * (alpha_t * (phi_1 / h + 1.)) * (model_s1 - model_s)
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_1 = torch.expm1(h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (
+                torch.exp(log_alpha_s1 - log_alpha_s) * x
+                - (sigma_s1 * phi_11) * model_s
+            )
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    torch.exp(log_alpha_t - log_alpha_s) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (0.5 / r1) * (sigma_t * phi_1) * (model_s1 - model_s)
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                    torch.exp(log_alpha_t - log_alpha_s) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1. / r1) * (sigma_t * (phi_1 / h - 1.)) * (model_s1 - model_s)
+                )
+        if return_intermediate:
+            return x_t, {'model_s': model_s, 'model_s1': model_s1}
+        else:
+            return x_t
+
+    def singlestep_dpm_solver_third_update(self, x, s, t, r1=1./3., r2=2./3., model_s=None, model_s1=None, return_intermediate=False, solver_type='dpmsolver'):
+        """
+        Singlestep solver DPM-Solver-3 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            r1: A `float`. The hyperparameter of the third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
+                If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpmsolver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type))
+        if r1 is None:
+            r1 = 1. / 3.
+        if r2 is None:
+            r2 = 2. / 3.
+        ns = self.noise_schedule
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        lambda_s2 = lambda_s + r2 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        s2 = ns.inverse_lambda(lambda_s2)
+        log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(s2), ns.marginal_std(t)
+        alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_11 = torch.expm1(-r1 * h)
+            phi_12 = torch.expm1(-r2 * h)
+            phi_1 = torch.expm1(-h)
+            phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
+            phi_2 = phi_1 / h + 1.
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (
+                    (sigma_s1 / sigma_s) * x
+                    - (alpha_s1 * phi_11) * model_s
+                )
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                (sigma_s2 / sigma_s) * x
+                - (alpha_s2 * phi_12) * model_s
+                + r2 / r1 * (alpha_s2 * phi_22) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1. / r2) * (alpha_t * phi_2) * (model_s2 - model_s)
+                )
+            elif solver_type == 'taylor':
+                D1_0 = (1. / r1) * (model_s1 - model_s)
+                D1_1 = (1. / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (alpha_t * phi_2) * D1
+                    - (alpha_t * phi_3) * D2
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_12 = torch.expm1(r2 * h)
+            phi_1 = torch.expm1(h)
+            phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
+            phi_2 = phi_1 / h - 1.
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (
+                    (torch.exp(log_alpha_s1 - log_alpha_s)) * x
+                    - (sigma_s1 * phi_11) * model_s
+                )
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                (torch.exp(log_alpha_s2 - log_alpha_s)) * x
+                - (sigma_s2 * phi_12) * model_s
+                - r2 / r1 * (sigma_s2 * phi_22) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_s)) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1. / r2) * (sigma_t * phi_2) * (model_s2 - model_s)
+                )
+            elif solver_type == 'taylor':
+                D1_0 = (1. / r1) * (model_s1 - model_s)
+                D1_1 = (1. / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_s)) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (sigma_t * phi_2) * D1
+                    - (sigma_t * phi_3) * D2
+                )
+
+        if return_intermediate:
+            return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
+        else:
+            return x_t
+
+    def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpmsolver"):
+        """
+        Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpmsolver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type))
+        ns = self.noise_schedule
+        model_prev_1, model_prev_0 = model_prev_list[-2], model_prev_list[-1]
+        t_prev_1, t_prev_0 = t_prev_list[-2], t_prev_list[-1]
+        lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0 = h_0 / h
+        D1_0 = (1. / r0) * (model_prev_0 - model_prev_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    (sigma_t / sigma_prev_0) * x
+                    - (alpha_t * phi_1) * model_prev_0
+                    - 0.5 * (alpha_t * phi_1) * D1_0
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                    (sigma_t / sigma_prev_0) * x
+                    - (alpha_t * phi_1) * model_prev_0
+                    + (alpha_t * (phi_1 / h + 1.)) * D1_0
+                )
+        else:
+            phi_1 = torch.expm1(h)
+            if solver_type == 'dpmsolver':
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                    - (sigma_t * phi_1) * model_prev_0
+                    - 0.5 * (sigma_t * phi_1) * D1_0
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                    - (sigma_t * phi_1) * model_prev_0
+                    - (sigma_t * (phi_1 / h - 1.)) * D1_0
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpmsolver'):
+        """
+        Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        model_prev_2, model_prev_1, model_prev_0 = model_prev_list
+        t_prev_2, t_prev_1, t_prev_0 = t_prev_list
+        lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_1 = lambda_prev_1 - lambda_prev_2
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0, r1 = h_0 / h, h_1 / h
+        D1_0 = (1. / r0) * (model_prev_0 - model_prev_1)
+        D1_1 = (1. / r1) * (model_prev_1 - model_prev_2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1. / (r0 + r1)) * (D1_0 - D1_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            phi_2 = phi_1 / h + 1.
+            phi_3 = phi_2 / h - 0.5
+            x_t = (
+                (sigma_t / sigma_prev_0) * x
+                - (alpha_t * phi_1) * model_prev_0
+                + (alpha_t * phi_2) * D1
+                - (alpha_t * phi_3) * D2
+            )
+        else:
+            phi_1 = torch.expm1(h)
+            phi_2 = phi_1 / h - 1.
+            phi_3 = phi_2 / h - 0.5
+            x_t = (
+                (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                - (sigma_t * phi_1) * model_prev_0
+                - (sigma_t * phi_2) * D1
+                - (sigma_t * phi_3) * D2
+            )
+        return x_t
+
+    def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpmsolver', r1=None, r2=None):
+        """
+        Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+            r1: A `float`. The hyperparameter of the second-order or third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
+        elif order == 2:
+            return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1)
+        elif order == 3:
+            return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1, r2=r2)
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpmsolver'):
+        """
+        Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
+        elif order == 2:
+            return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+        elif order == 3:
+            return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5, solver_type='dpmsolver'):
+        """
+        The adaptive step size solver based on singlestep DPM-Solver.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `t_T`.
+            order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            h_init: A `float`. The initial step size (for logSNR).
+            atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
+            rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
+            theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
+            t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the 
+                current time and `t_0` is less than `t_err`. The default setting is 1e-5.
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_0: A pytorch tensor. The approximated solution at time `t_0`.
+
+        [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
+        """
+        ns = self.noise_schedule
+        s = t_T * torch.ones((1,)).to(x)
+        lambda_s = ns.marginal_lambda(s)
+        lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
+        h = h_init * torch.ones_like(s).to(x)
+        x_prev = x
+        nfe = 0
+        if order == 2:
+            r1 = 0.5
+            lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
+            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, solver_type=solver_type, **kwargs)
+        elif order == 3:
+            r1, r2 = 1. / 3., 2. / 3.
+            lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, return_intermediate=True, solver_type=solver_type)
+            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2, solver_type=solver_type, **kwargs)
+        else:
+            raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
+        while torch.abs((s - t_0)).mean() > t_err:
+            t = ns.inverse_lambda(lambda_s + h)
+            x_lower, lower_noise_kwargs = lower_update(x, s, t)
+            x_higher = higher_update(x, s, t, **lower_noise_kwargs)
+            delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
+            norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
+            E = norm_fn((x_higher - x_lower) / delta).max()
+            if torch.all(E <= 1.):
+                x = x_higher
+                s = t
+                x_prev = x_lower
+                lambda_s = ns.marginal_lambda(s)
+            h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
+            nfe += order
+        print('adaptive solver nfe', nfe)
+        return x
+
+    def add_noise(self, x, t, noise=None):
+        """
+        Compute the noised input xt = alpha_t * x + sigma_t * noise. 
+
+        Args:
+            x: A `torch.Tensor` with shape `(batch_size, *shape)`.
+            t: A `torch.Tensor` with shape `(t_size,)`.
+        Returns:
+            xt with shape `(t_size, batch_size, *shape)`.
+        """
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
+        if noise is None:
+            noise = torch.randn((t.shape[0], *x.shape), device=x.device)
+        x = x.reshape((-1, *x.shape))
+        xt = expand_dims(alpha_t, x.dim()) * x + expand_dims(sigma_t, x.dim()) * noise
+        if t.shape[0] == 1:
+            return xt.squeeze(0)
+        else:
+            return xt
+
+    def inverse(self, x, steps=20, t_start=None, t_end=None, order=2, skip_type='time_uniform',
+        method='multistep', lower_order_final=True, denoise_to_zero=False, solver_type='dpmsolver',
+        atol=0.0078, rtol=0.05, return_intermediate=False,
+    ):
+        """
+        Inverse the sample `x` from time `t_start` to `t_end` by DPM-Solver.
+        For discrete-time DPMs, we use `t_start=1/N`, where `N` is the total time steps during training.
+        """
+        t_0 = 1. / self.noise_schedule.total_N if t_start is None else t_start
+        t_T = self.noise_schedule.T if t_end is None else t_end
+        assert t_0 > 0 and t_T > 0, "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
+        return self.sample(x, steps=steps, t_start=t_0, t_end=t_T, order=order, skip_type=skip_type,
+            method=method, lower_order_final=lower_order_final, denoise_to_zero=denoise_to_zero, solver_type=solver_type,
+            atol=atol, rtol=rtol, return_intermediate=return_intermediate)
+
+    def sample(self, x, steps=20, t_start=None, t_end=None, order=2, skip_type='time_uniform',
+        method='multistep', lower_order_final=True, denoise_to_zero=False, solver_type='dpmsolver',
+        atol=0.0078, rtol=0.05, return_intermediate=False,
+    ):
+        """
+        Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
+
+        =====================================================
+
+        We support the following algorithms for both noise prediction model and data prediction model:
+            - 'singlestep':
+                Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver. 
+                We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
+                The total number of function evaluations (NFE) == `steps`.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    - If `order` == 1:
+                        - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
+                        - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
+                        - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                    - If `order` == 3:
+                        - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                        - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
+            - 'multistep':
+                Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
+                We initialize the first `order` values by lower order multistep solvers.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    Denote K = steps.
+                    - If `order` == 1:
+                        - We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
+                    - If `order` == 3:
+                        - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
+            - 'singlestep_fixed':
+                Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
+                We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
+            - 'adaptive':
+                Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
+                We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
+                You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
+                (NFE) and the sample quality.
+                    - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
+                    - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
+
+        =====================================================
+
+        Some advices for choosing the algorithm:
+            - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
+                Use singlestep DPM-Solver or DPM-Solver++ ("DPM-Solver-fast" in the paper) with `order = 3`.
+                e.g., DPM-Solver:
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+                            skip_type='time_uniform', method='singlestep')
+                e.g., DPM-Solver++:
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+                            skip_type='time_uniform', method='singlestep')
+            - For **guided sampling with large guidance scale** by DPMs:
+                Use multistep DPM-Solver with `algorithm_type="dpmsolver++"` and `order = 2`.
+                e.g.
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
+                            skip_type='time_uniform', method='multistep')
+
+        We support three types of `skip_type`:
+            - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
+            - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
+            - 'time_quadratic': quadratic time for the time steps.
+
+        =====================================================
+        Args:
+            x: A pytorch tensor. The initial value at time `t_start`
+                e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
+            steps: A `int`. The total number of function evaluations (NFE).
+            t_start: A `float`. The starting time of the sampling.
+                If `T` is None, we use self.noise_schedule.T (default is 1.0).
+            t_end: A `float`. The ending time of the sampling.
+                If `t_end` is None, we use 1. / self.noise_schedule.total_N.
+                e.g. if total_N == 1000, we have `t_end` == 1e-3.
+                For discrete-time DPMs:
+                    - We recommend `t_end` == 1. / self.noise_schedule.total_N.
+                For continuous-time DPMs:
+                    - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
+            order: A `int`. The order of DPM-Solver.
+            skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
+            method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
+            denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
+                Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
+
+                This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
+                score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
+                for diffusion models sampling by diffusion SDEs for low-resolutional images
+                (such as CIFAR-10). However, we observed that such trick does not matter for
+                high-resolutional images. As it needs an additional NFE, we do not recommend
+                it for high-resolutional images.
+            lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
+                Only valid for `method=multistep` and `steps < 15`. We empirically find that
+                this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
+                (especially for steps <= 10). So we recommend to set it to be `True`.
+            solver_type: A `str`. The taylor expansion type for the solver. `dpmsolver` or `taylor`. We recommend `dpmsolver`.
+            atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+            rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+            return_intermediate: A `bool`. Whether to save the xt at each step.
+                When set to `True`, method returns a tuple (x0, intermediates); when set to False, method returns only x0.
+        Returns:
+            x_end: A pytorch tensor. The approximated solution at time `t_end`.
+
+        """
+        t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
+        t_T = self.noise_schedule.T if t_start is None else t_start
+        assert t_0 > 0 and t_T > 0, "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
+        if return_intermediate:
+            assert method in ['multistep', 'singlestep', 'singlestep_fixed'], "Cannot use adaptive solver when saving intermediate values"
+        if self.correcting_xt_fn is not None:
+            assert method in ['multistep', 'singlestep', 'singlestep_fixed'], "Cannot use adaptive solver when correcting_xt_fn is not None"
+        device = x.device
+        intermediates = []
+        with torch.no_grad():
+            if method == 'adaptive':
+                x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol, solver_type=solver_type)
+            elif method == 'multistep':
+                assert steps >= order
+                timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
+                assert timesteps.shape[0] - 1 == steps
+                # Init the initial values.
+                step = 0
+                t = timesteps[step]
+                t_prev_list = [t]
+                model_prev_list = [self.model_fn(x, t)]
+                if self.correcting_xt_fn is not None:
+                    x = self.correcting_xt_fn(x, t, step)
+                if return_intermediate:
+                    intermediates.append(x)
+                # Init the first `order` values by lower order multistep DPM-Solver.
+                for step in range(1, order):
+                    t = timesteps[step]
+                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, step, solver_type=solver_type)
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+                    t_prev_list.append(t)
+                    model_prev_list.append(self.model_fn(x, t))
+                # Compute the remaining values by `order`-th order multistep DPM-Solver.
+                for step in range(order, steps + 1):
+                    t = timesteps[step]
+                    # We only use lower order for steps < 10
+                    if lower_order_final and steps < 10:
+                        step_order = min(order, steps + 1 - step)
+                    else:
+                        step_order = order
+                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, step_order, solver_type=solver_type)
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+                    for i in range(order - 1):
+                        t_prev_list[i] = t_prev_list[i + 1]
+                        model_prev_list[i] = model_prev_list[i + 1]
+                    t_prev_list[-1] = t
+                    # We do not need to evaluate the final model value.
+                    if step < steps:
+                        model_prev_list[-1] = self.model_fn(x, t)
+            elif method in ['singlestep', 'singlestep_fixed']:
+                if method == 'singlestep':
+                    timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order, skip_type=skip_type, t_T=t_T, t_0=t_0, device=device)
+                elif method == 'singlestep_fixed':
+                    K = steps // order
+                    orders = [order,] * K
+                    timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
+                for step, order in enumerate(orders):
+                    s, t = timesteps_outer[step], timesteps_outer[step + 1]
+                    timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=s.item(), t_0=t.item(), N=order, device=device)
+                    lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
+                    h = lambda_inner[-1] - lambda_inner[0]
+                    r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
+                    r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
+                    x = self.singlestep_dpm_solver_update(x, s, t, order, solver_type=solver_type, r1=r1, r2=r2)
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+            else:
+                raise ValueError("Got wrong method {}".format(method))
+            if denoise_to_zero:
+                t = torch.ones((1,)).to(device) * t_0
+                x = self.denoise_to_zero_fn(x, t)
+                if self.correcting_xt_fn is not None:
+                    x = self.correcting_xt_fn(x, t, step + 1)
+                if return_intermediate:
+                    intermediates.append(x)
+        if return_intermediate:
+            return x, intermediates
+        else:
+            return x
+
+
+
+#############################################################
+# other utility functions
+#############################################################
+
+def interpolate_fn(x, xp, yp):
+    """
+    A piecewise linear function y = f(x), using xp and yp as keypoints.
+    We implement f(x) in a differentiable way (i.e. applicable for autograd).
+    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
+
+    Args:
+        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
+        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
+        yp: PyTorch tensor with shape [C, K].
+    Returns:
+        The function values f(x), with shape [N, C].
+    """
+    N, K = x.shape[0], xp.shape[1]
+    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
+    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
+    x_idx = torch.argmin(x_indices, dim=2)
+    cand_start_idx = x_idx - 1
+    start_idx = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(1, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
+    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
+    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
+    start_idx2 = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(0, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
+    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
+    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
+    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
+    return cand
+
+
+def expand_dims(v, dims):
+    """
+    Expand the tensor `v` to the dim `dims`.
+
+    Args:
+        `v`: a PyTorch tensor with shape [N].
+        `dim`: a `int`.
+    Returns:
+        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
+    """
+    return v[(...,) + (None,)*(dims - 1)]

src/dpm_solver/pipeline_dpm_solver.py

@@ -0,0 +1,117 @@
+import torch
+from PIL import Image
+
+from .dpm_solver_pytorch import (NoiseScheduleVP, 
+                                model_wrapper, 
+                                DPM_Solver)
+
+class FontDiffuserDPMPipeline():
+    """FontDiffuser pipeline with DPM_Solver scheduler.
+    """
+    
+    def __init__(
+        self, 
+        model, 
+        ddpm_train_scheduler,
+        version="V3",
+        model_type="noise",
+        guidance_type="classifier-free",
+        guidance_scale=7.5
+    ):
+        super().__init__()
+        self.model = model
+        self.train_scheduler_betas = ddpm_train_scheduler.betas
+        # Define the noise schedule
+        self.noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.train_scheduler_betas)
+
+        self.version = version
+        self.model_type = model_type
+        self.guidance_type = guidance_type
+        self.guidance_scale = guidance_scale
+
+    def numpy_to_pil(self, images):
+        """Convert a numpy image or a batch of images to a PIL image.
+        """
+        if images.ndim == 3:
+            images = images[None, ...]
+        images = (images * 255).round().astype("uint8")
+        pil_images = [Image.fromarray(image) for image in images]
+
+        return pil_images
+
+    def generate(
+        self,
+        content_images,
+        style_images,
+        batch_size,
+        order,
+        num_inference_step,
+        content_encoder_downsample_size,
+        t_start=None,
+        t_end=None,
+        dm_size=(96, 96),
+        algorithm_type="dpmsolver++",
+        skip_type="time_uniform",
+        method="multistep",
+        correcting_x0_fn=None,
+        generator=None,
+    ):
+        model_kwargs = {}
+        model_kwargs["version"] = self.version
+        model_kwargs["content_encoder_downsample_size"] = content_encoder_downsample_size
+
+        cond = []
+        cond.append(content_images)
+        cond.append(style_images)
+
+        uncond = []
+        uncond_content_images = torch.ones_like(content_images).to(self.model.device)
+        uncond_style_images = torch.ones_like(style_images).to(self.model.device)
+        uncond.append(uncond_content_images)
+        uncond.append(uncond_style_images)
+
+        # 2.Convert the discrete-time model to the continuous-time
+        model_fn = model_wrapper(
+            model=self.model,
+            noise_schedule=self.noise_schedule,
+            model_type=self.model_type,
+            model_kwargs=model_kwargs,
+            guidance_type=self.guidance_type,
+            condition=cond, 
+            unconditional_condition=uncond,
+            guidance_scale=self.guidance_scale
+        )
+
+        # 3. Define dpm-solver and sample by multistep DPM-Solver.
+        # (We recommend multistep DPM-Solver for conditional sampling)
+        # You can adjust the `steps` to balance the computation costs and the sample quality.
+        dpm_solver = DPM_Solver(
+            model_fn=model_fn,
+            noise_schedule=self.noise_schedule,
+            algorithm_type=algorithm_type,
+            correcting_x0_fn=correcting_x0_fn
+        )
+        # If the DPM is defined on pixel-space images, you can further set `correcting_x0_fn="dynamic_thresholding"
+
+        # 4. Generate
+        # Sample gaussian noise to begin loop => [batch, 3, height, width]
+        x_T = torch.randn(
+            (batch_size, 3, dm_size[0], dm_size[1]),
+            generator=generator,
+        )
+        x_T = x_T.to(self.model.device)
+
+        x_sample = dpm_solver.sample(
+            x=x_T,
+            steps=num_inference_step,
+            order=order,
+            skip_type=skip_type,
+            method=method,
+        )
+
+        x_sample = (x_sample / 2 + 0.5).clamp(0, 1)
+        x_sample = x_sample.cpu().permute(0, 2, 3, 1).numpy()
+    
+        x_images = self.numpy_to_pil(x_sample)
+
+        return x_images

src/model.py

@@ -0,0 +1,110 @@
+import math
+import torch
+import torch.nn as nn
+
+from diffusers import ModelMixin
+from diffusers.configuration_utils import (ConfigMixin, 
+                                           register_to_config)
+
+class FontDiffuserModel(ModelMixin, ConfigMixin):
+    """Forward function for FontDiffuer with content encoder \
+        style encoder and unet.
+    """
+
+    @register_to_config
+    def __init__(
+        self, 
+        unet, 
+        style_encoder,
+        content_encoder,
+    ):
+        super().__init__()
+        self.unet = unet
+        self.style_encoder = style_encoder
+        self.content_encoder = content_encoder
+    
+    def forward(
+        self, 
+        x_t, 
+        timesteps, 
+        style_images,
+        content_images,
+        content_encoder_downsample_size,
+    ):
+        style_img_feature, _, _ = self.style_encoder(style_images)
+    
+        batch_size, channel, height, width = style_img_feature.shape
+        style_hidden_states = style_img_feature.permute(0, 2, 3, 1).reshape(batch_size, height*width, channel)
+    
+        # Get the content feature
+        content_img_feature, content_residual_features = self.content_encoder(content_images)
+        content_residual_features.append(content_img_feature)
+        # Get the content feature from reference image
+        style_content_feature, style_content_res_features = self.content_encoder(style_images)
+        style_content_res_features.append(style_content_feature)
+
+        input_hidden_states = [style_img_feature, content_residual_features, \
+                               style_hidden_states, style_content_res_features]
+
+        out = self.unet(
+            x_t, 
+            timesteps, 
+            encoder_hidden_states=input_hidden_states,
+            content_encoder_downsample_size=content_encoder_downsample_size,
+        )
+        noise_pred = out[0]
+        offset_out_sum = out[1]
+        
+        return noise_pred, offset_out_sum
+
+
+class FontDiffuserModelDPM(ModelMixin, ConfigMixin):
+    """DPM Forward function for FontDiffuer with content encoder \
+        style encoder and unet.
+    """
+    @register_to_config
+    def __init__(
+        self, 
+        unet, 
+        style_encoder,
+        content_encoder,
+    ):
+        super().__init__()
+        self.unet = unet
+        self.style_encoder = style_encoder
+        self.content_encoder = content_encoder
+    
+    def forward(
+        self, 
+        x_t, 
+        timesteps, 
+        cond,
+        content_encoder_downsample_size,
+        version,
+    ):
+        content_images = cond[0]
+        style_images = cond[1]
+
+        style_img_feature, _, style_residual_features = self.style_encoder(style_images)
+        
+        batch_size, channel, height, width = style_img_feature.shape
+        style_hidden_states = style_img_feature.permute(0, 2, 3, 1).reshape(batch_size, height*width, channel)
+        
+        # Get content feature
+        content_img_feture, content_residual_features = self.content_encoder(content_images)
+        content_residual_features.append(content_img_feture)
+        # Get the content feature from reference image
+        style_content_feature, style_content_res_features = self.content_encoder(style_images)
+        style_content_res_features.append(style_content_feature)
+
+        input_hidden_states = [style_img_feature, content_residual_features, style_hidden_states, style_content_res_features]
+
+        out = self.unet(
+            x_t, 
+            timesteps, 
+            encoder_hidden_states=input_hidden_states,
+            content_encoder_downsample_size=content_encoder_downsample_size,
+        )
+        noise_pred = out[0]
+        
+        return noise_pred

src/modules/__init__.py

@@ -0,0 +1,3 @@
+from .content_encoder import ContentEncoder
+from .style_encoder import StyleEncoder
+from .unet import UNet

src/modules/attention.py

@@ -0,0 +1,414 @@
+from typing import Optional
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data. First, project the input (aka embedding) and reshape to b, t, d. Then apply
+    standard transformer action. Finally, reshape to image.
+
+    Parameters:
+        in_channels (:obj:`int`): The number of channels in the input and output.
+        n_heads (:obj:`int`): The number of heads to use for multi-head attention.
+        d_head (:obj:`int`): The number of channels in each head.
+        depth (:obj:`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+        dropout (:obj:`float`, *optional*, defaults to 0.1): The dropout probability to use.
+        context_dim (:obj:`int`, *optional*): The number of context dimensions to use.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        n_heads: int,
+        d_head: int,
+        depth: int = 1,
+        dropout: float = 0.0,
+        num_groups: int = 32,
+        context_dim: Optional[int] = None,
+    ):
+        super().__init__()
+        self.n_heads = n_heads
+        self.d_head = d_head
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
+                for d in range(depth)
+            ]
+        )
+
+        self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def _set_attention_slice(self, slice_size):
+        for block in self.transformer_blocks:
+            block._set_attention_slice(slice_size)
+
+    def forward(self, hidden_states, context=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        batch, channel, height, weight = hidden_states.shape
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.proj_in(hidden_states)
+        inner_dim = hidden_states.shape[1]
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)  # here change the shape torch.Size([1, 4096, 128])
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, context=context)  # hidden_states: torch.Size([1, 4096, 128])
+        hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2)   # torch.Size([1, 128, 64, 64])
+        hidden_states = self.proj_out(hidden_states)
+        return hidden_states + residual
+
+
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (:obj:`int`): The number of channels in the input and output.
+        n_heads (:obj:`int`): The number of heads to use for multi-head attention.
+        d_head (:obj:`int`): The number of channels in each head.
+        dropout (:obj:`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        context_dim (:obj:`int`, *optional*): The size of the context vector for cross attention.
+        gated_ff (:obj:`bool`, *optional*, defaults to :obj:`False`): Whether to use a gated feed-forward network.
+        checkpoint (:obj:`bool`, *optional*, defaults to :obj:`False`): Whether to use checkpointing.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        d_head: int,
+        dropout=0.0,
+        context_dim: Optional[int] = None,
+        gated_ff: bool = True,
+        checkpoint: bool = True,
+    ):
+        super().__init__()
+        self.attn1 = CrossAttention(
+            query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout
+        )  # is a self-attention
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = CrossAttention(
+            query_dim=dim, context_dim=context_dim, heads=n_heads, dim_head=d_head, dropout=dropout
+        )  # is self-attn if context is none
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def _set_attention_slice(self, slice_size):
+        self.attn1._slice_size = slice_size
+        self.attn2._slice_size = slice_size
+
+    def forward(self, hidden_states, context=None):
+        hidden_states = hidden_states.contiguous() if hidden_states.device.type == "mps" else hidden_states
+        hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states   # hidden_states: torch.Size([1, 4096, 128])
+        hidden_states = self.attn2(self.norm2(hidden_states), context=context) + hidden_states
+        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+        return hidden_states
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (:obj:`int`): The number of channels in the input.
+        dim_out (:obj:`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (:obj:`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        glu (:obj:`bool`, *optional*, defaults to :obj:`False`): Whether to use GLU activation.
+        dropout (:obj:`float`, *optional*, defaults to 0.0): The dropout probability to use.
+    """
+
+    def __init__(
+        self, dim: int, dim_out: Optional[int] = None, mult: int = 4, glu: bool = False, dropout: float = 0.0
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        project_in = GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, hidden_states):
+        return self.net(hidden_states)
+
+
+class GEGLU(nn.Module):
+    r"""
+    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
+
+    Parameters:
+        dim_in (:obj:`int`): The number of channels in the input.
+        dim_out (:obj:`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, hidden_states):
+        hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
+        return hidden_states * F.gelu(gate)
+
+
+class CrossAttention(nn.Module):
+    r"""
+    A cross attention layer.
+
+    Parameters:
+        query_dim (:obj:`int`): The number of channels in the query.
+        context_dim (:obj:`int`, *optional*):
+            The number of channels in the context. If not given, defaults to `query_dim`.
+        heads (:obj:`int`,  *optional*, defaults to 8): The number of heads to use for multi-head attention.
+        dim_head (:obj:`int`,  *optional*, defaults to 64): The number of channels in each head.
+        dropout (:obj:`float`, *optional*, defaults to 0.0): The dropout probability to use.
+    """
+
+    def __init__(
+        self, query_dim: int, context_dim: Optional[int] = None, heads: int = 8, dim_head: int = 64, dropout: int = 0.0
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = context_dim if context_dim is not None else query_dim
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set slice_size with `set_attention_slice`
+        self._slice_size = None
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
+
+    def reshape_heads_to_batch_dim(self, tensor):
+        batch_size, seq_len, dim = tensor.shape
+        head_size = self.heads
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)
+        return tensor
+
+    def reshape_batch_dim_to_heads(self, tensor):
+        batch_size, seq_len, dim = tensor.shape
+        head_size = self.heads
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        return tensor
+
+    def forward(self, hidden_states, context=None, mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        query = self.to_q(hidden_states)
+        context = context if context is not None else hidden_states
+        key = self.to_k(context)
+        value = self.to_v(context)
+
+        dim = query.shape[-1]
+
+        query = self.reshape_heads_to_batch_dim(query)
+        key = self.reshape_heads_to_batch_dim(key)
+        value = self.reshape_heads_to_batch_dim(value)
+
+        # TODO(PVP) - mask is currently never used. Remember to re-implement when used
+
+        # attention, what we cannot get enough of
+
+        if self._slice_size is None or query.shape[0] // self._slice_size == 1:
+            hidden_states = self._attention(query, key, value)
+        else:
+            hidden_states = self._sliced_attention(query, key, value, sequence_length, dim)
+
+        return self.to_out(hidden_states)
+
+    def _attention(self, query, key, value):
+        # TODO: use baddbmm for better performance
+        attention_scores = torch.matmul(query, key.transpose(-1, -2)) * self.scale
+        attention_probs = attention_scores.softmax(dim=-1)
+        # compute attention output
+        hidden_states = torch.matmul(attention_probs, value)
+        # reshape hidden_states
+        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
+        return hidden_states
+
+    def _sliced_attention(self, query, key, value, sequence_length, dim):
+        batch_size_attention = query.shape[0]
+        hidden_states = torch.zeros(
+            (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype
+        )
+        slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]
+        for i in range(hidden_states.shape[0] // slice_size):
+            start_idx = i * slice_size
+            end_idx = (i + 1) * slice_size
+            attn_slice = (
+                torch.matmul(query[start_idx:end_idx], key[start_idx:end_idx].transpose(1, 2)) * self.scale
+            )  # TODO: use baddbmm for better performance
+            attn_slice = attn_slice.softmax(dim=-1)
+            attn_slice = torch.matmul(attn_slice, value[start_idx:end_idx])
+
+            hidden_states[start_idx:end_idx] = attn_slice
+
+        # reshape hidden_states
+        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
+        return hidden_states
+
+
+class OffsetRefStrucInter(nn.Module):
+
+    def __init__(
+        self,
+        res_in_channels: int,
+        style_feat_in_channels: int,
+        n_heads: int,
+        num_groups: int = 32,
+        dropout: float = 0.0,
+        gated_ff: bool = True,
+    ):  
+        super().__init__()
+        # style feat projecter
+        self.style_proj_in = nn.Conv2d(style_feat_in_channels, style_feat_in_channels, kernel_size=1, stride=1, padding=0)
+        self.gnorm_s = torch.nn.GroupNorm(num_groups=num_groups, num_channels=style_feat_in_channels, eps=1e-6, affine=True)
+        self.ln_s = nn.LayerNorm(style_feat_in_channels)
+
+        # content feat projecter
+        self.content_proj_in = nn.Conv2d(res_in_channels, res_in_channels, kernel_size=1, stride=1, padding=0)
+        self.gnorm_c = torch.nn.GroupNorm(num_groups=num_groups, num_channels=res_in_channels, eps=1e-6, affine=True)
+        self.ln_c = nn.LayerNorm(res_in_channels)
+
+        # cross-attention
+        # dim_head is the middle dealing dimension, output dimension will be change to quert_dim by Linear
+        self.cross_attention = CrossAttention(
+            query_dim=style_feat_in_channels, context_dim=res_in_channels, heads=n_heads, dim_head=res_in_channels, dropout=dropout
+        )
+
+        # FFN
+        self.ff = FeedForward(style_feat_in_channels, dropout=dropout, glu=gated_ff)
+        self.ln_ff = nn.LayerNorm(style_feat_in_channels)
+
+        self.gnorm_out = torch.nn.GroupNorm(num_groups=num_groups, num_channels=style_feat_in_channels, eps=1e-6, affine=True)
+        self.proj_out = nn.Conv2d(style_feat_in_channels, 1*2*3*3, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, res_hidden_states, style_content_hidden_states):
+        batch, c_channel, height, width = res_hidden_states.shape
+        _, s_channel, _, _ = style_content_hidden_states.shape
+        # style projecter
+        style_content_hidden_states = self.gnorm_s(style_content_hidden_states)
+        style_content_hidden_states = self.style_proj_in(style_content_hidden_states)
+        
+        style_content_hidden_states = style_content_hidden_states.permute(0, 2, 3, 1).reshape(batch, height*width, s_channel)
+        style_content_hidden_states = self.ln_s(style_content_hidden_states)
+
+        # content projecter
+        res_hidden_states = self.gnorm_c(res_hidden_states)
+        res_hidden_states = self.content_proj_in(res_hidden_states)
+        
+        res_hidden_states = res_hidden_states.permute(0, 2, 3, 1).reshape(batch, height*width, c_channel)
+        res_hidden_states = self.ln_c(res_hidden_states)
+
+        # style and content cross-attention
+        hidden_states = self.cross_attention(style_content_hidden_states, context=res_hidden_states)
+
+        # ffn
+        hidden_states = self.ff(self.ln_ff(hidden_states)) + hidden_states
+
+        # reshape
+        _, _, c = hidden_states.shape
+        reshape_out = hidden_states.permute(0, 2, 1).reshape(batch, c, height, width)
+
+        # projert out
+        reshape_out = self.gnorm_out(reshape_out)
+        offset_out = self.proj_out(reshape_out)
+
+        return offset_out
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            # nn.ReLU(inplace=True),
+            nn.SiLU(),
+            nn.Linear(channel // reduction, channel, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+
+class Mish(torch.nn.Module):
+    def forward(self, hidden_states):
+        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))
+
+
+class ChannelAttnBlock(nn.Module):
+    """This is the Channel Attention in MCA.
+    """
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        groups=32,
+        groups_out=None,
+        eps=1e-6,
+        non_linearity="swish",
+        channel_attn=False,
+        reduction=32):
+        super().__init__()
+
+        if groups_out is None:
+            groups_out = groups
+
+        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1)
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+        
+        self.channel_attn = channel_attn
+        if self.channel_attn:
+            # SE Attention
+            self.se_channel_attn = SELayer(channel=in_channels, reduction=reduction)
+
+        # Down channel: Use the conv1*1 to down the channel wise
+        self.norm3 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        self.down_channel = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1) # conv1*1
+
+    def forward(self, input, content_feature):
+
+        concat_feature = torch.cat([input, content_feature], dim=1)
+        hidden_states = concat_feature
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        if self.channel_attn:
+            hidden_states = self.se_channel_attn(hidden_states)
+            hidden_states = hidden_states + concat_feature
+
+        # Down channel
+        hidden_states = self.norm3(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.down_channel(hidden_states)
+
+        return hidden_states

src/modules/content_encoder.py

@@ -0,0 +1,435 @@
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+from torch.nn import Parameter as P
+
+from diffusers import ModelMixin
+from diffusers.configuration_utils import (ConfigMixin, 
+                                           register_to_config)
+
+
+def proj(x, y):
+   return torch.mm(y, x.t()) * y / torch.mm(y, y.t())
+
+
+def gram_schmidt(x, ys):
+    for y in ys:
+        x = x - proj(x, y)
+    return x
+
+
+def power_iteration(W, u_, update=True, eps=1e-12):
+    us, vs, svs = [], [], []
+    for i, u in enumerate(u_):
+        with torch.no_grad():
+            v = torch.matmul(u, W)
+            v = F.normalize(gram_schmidt(v, vs), eps=eps)
+            vs += [v]
+            u = torch.matmul(v, W.t())
+            u = F.normalize(gram_schmidt(u, us), eps=eps)
+            us += [u]
+            if update:
+                u_[i][:] = u
+        svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]
+    return svs, us, vs
+
+
+class LinearBlock(nn.Module):
+    def __init__(
+        self, 
+        in_dim, 
+        out_dim, 
+        norm='none', 
+        act='relu', 
+        use_sn=False
+    ):
+        super(LinearBlock, self).__init__()
+        use_bias = True
+        self.fc = nn.Linear(in_dim, out_dim, bias=use_bias)
+        if use_sn:
+            self.fc = nn.utils.spectral_norm(self.fc)
+
+        # initialize normalization
+        norm_dim = out_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm1d(norm_dim)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm1d(norm_dim)
+        elif norm == 'none':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if act == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif act == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif act == 'tanh':
+            self.activation = nn.Tanh()
+        elif act == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(act)
+
+    def forward(self, x):
+        out = self.fc(x)
+        if self.norm:
+            out = self.norm(out)
+        if self.activation:
+            out = self.activation(out)
+        return out
+
+
+class MLP(nn.Module):
+    def __init__(
+        self, 
+        nf_in, 
+        nf_out, 
+        nf_mlp, 
+        num_blocks, 
+        norm, 
+        act, 
+        use_sn =False
+    ):
+        super(MLP,self).__init__()
+        self.model = nn.ModuleList()
+        nf = nf_mlp
+        self.model.append(LinearBlock(nf_in, nf, norm = norm, act = act, use_sn = use_sn))
+        for _ in range((num_blocks - 2)):
+            self.model.append(LinearBlock(nf, nf, norm=norm, act=act, use_sn=use_sn))
+        self.model.append(LinearBlock(nf, nf_out, norm='none', act ='none', use_sn = use_sn))
+        self.model = nn.Sequential(*self.model)
+    
+    def forward(self, x):
+        return self.model(x.view(x.size(0), -1))
+
+
+class SN(object):
+    def __init__(
+        self, 
+        num_svs, 
+        num_itrs, 
+        num_outputs, 
+        transpose=False, 
+        eps=1e-12
+    ):
+        self.num_itrs = num_itrs
+        self.num_svs = num_svs
+        self.transpose = transpose
+        self.eps = eps
+        for i in range(self.num_svs):
+            self.register_buffer('u%d' % i, torch.randn(1, num_outputs))
+            self.register_buffer('sv%d' % i, torch.ones(1))
+
+    @property
+    def u(self):
+        return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]
+
+    @property
+    def sv(self):
+        return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]
+
+    def W_(self):
+        W_mat = self.weight.view(self.weight.size(0), -1)
+        if self.transpose:
+            W_mat = W_mat.t()
+        for _ in range(self.num_itrs):
+            svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)
+        if self.training:
+            with torch.no_grad():
+                for i, sv in enumerate(svs):
+                    self.sv[i][:] = sv
+        return self.weight / svs[0]
+
+class SNConv2d(nn.Conv2d, SN):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                padding=0, dilation=1, groups=1, bias=True,
+                num_svs=1, num_itrs=1, eps=1e-12):
+        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride,
+                        padding, dilation, groups, bias)
+        SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)
+
+    def forward(self, x):
+        return F.conv2d(x, self.W_(), self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+    def forward_wo_sn(self, x):
+        return F.conv2d(x, self.weight, self.bias, self.stride,
+                    self.padding, self.dilation, self.groups)
+
+
+class SNLinear(nn.Linear, SN):
+    def __init__(self, in_features, out_features, bias=True,
+                num_svs=1, num_itrs=1, eps=1e-12):
+        nn.Linear.__init__(self, in_features, out_features, bias)
+        SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)
+
+    def forward(self, x):
+        return F.linear(x, self.W_(), self.bias)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self, 
+        ch, 
+        which_conv=SNConv2d, 
+        name='attention'
+    ):
+        super(Attention, self).__init__()
+        self.ch = ch
+        self.which_conv = which_conv
+        self.theta = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
+        self.phi = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
+        self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)
+        self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)
+        # Learnable gain parameter
+        self.gamma = P(torch.tensor(0.), requires_grad=True)
+
+    def forward(self, x, y=None):
+        theta = self.theta(x)
+        phi = F.max_pool2d(self.phi(x), [2,2])
+        g = F.max_pool2d(self.g(x), [2,2])
+        
+        theta = theta.view(-1, self. ch // 8, x.shape[2] * x.shape[3])
+        phi = phi.view(-1, self. ch // 8, x.shape[2] * x.shape[3] // 4)
+        g = g.view(-1, self. ch // 2, x.shape[2] * x.shape[3] // 4)
+        
+        beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)
+        
+        o = self.o(torch.bmm(g, beta.transpose(1,2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))
+        return self.gamma * o + x
+
+
+class DBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,
+                preactivation=False, activation=None, downsample=None,):
+        super(DBlock, self).__init__()
+        
+        self.in_channels, self.out_channels = in_channels, out_channels
+
+        self.hidden_channels = self.out_channels if wide else self.in_channels
+        self.which_conv = which_conv
+        self.preactivation = preactivation
+        self.activation = activation
+        self.downsample = downsample
+
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)
+        self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)
+        self.learnable_sc = True if (in_channels != out_channels) or downsample else False
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                            kernel_size=1, padding=0)
+    def shortcut(self, x):
+        if self.preactivation:
+            if self.learnable_sc:
+                x = self.conv_sc(x)
+            if self.downsample:
+                x = self.downsample(x)
+        else:
+            if self.downsample:
+                x = self.downsample(x)
+            if self.learnable_sc:
+                x = self.conv_sc(x)
+        return x
+
+    def forward(self, x):
+        if self.preactivation:
+            h = F.relu(x)
+        else:
+            h = x
+        h = self.conv1(h)
+        h = self.conv2(self.activation(h))
+        if self.downsample:
+            h = self.downsample(h)
+
+        return h + self.shortcut(x)
+
+
+class GBlock(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 which_conv=nn.Conv2d,which_bn= nn.BatchNorm2d, activation=None,
+                 upsample=None):
+        super(GBlock, self).__init__()
+
+        self.in_channels, self.out_channels = in_channels, out_channels
+        self.which_conv,self.which_bn =which_conv, which_bn
+        self.activation = activation
+        self.upsample = upsample
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.out_channels)
+        self.conv2 = self.which_conv(self.out_channels, self.out_channels)
+        self.learnable_sc = in_channels != out_channels or upsample
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                           kernel_size=1, padding=0)
+        # Batchnorm layers
+        self.bn1 = self.which_bn(in_channels)
+        self.bn2 = self.which_bn(out_channels)
+        # upsample layers
+        self.upsample = upsample
+
+    
+    def forward(self, x):
+        h = self.activation(self.bn1(x))
+        if self.upsample:
+            h = self.upsample(h)
+            x = self.upsample(x)
+        h = self.conv1(h)
+        h = self.activation(self.bn2(h))
+        h = self.conv2(h)
+        if self.learnable_sc:
+            x = self.conv_sc(x)
+        return h + x
+
+
+class GBlock2(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                which_conv=nn.Conv2d, activation=None,
+                upsample=None, skip_connection = True):
+        super(GBlock2, self).__init__()
+
+        self.in_channels, self.out_channels = in_channels, out_channels
+        self.which_conv = which_conv
+        self.activation = activation
+        self.upsample = upsample
+
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.out_channels)
+        self.conv2 = self.which_conv(self.out_channels, self.out_channels)
+        self.learnable_sc = in_channels != out_channels or upsample
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                            kernel_size=1, padding=0)
+
+        # upsample layers
+        self.upsample = upsample
+        self.skip_connection = skip_connection
+
+    def forward(self, x):
+        h = self.activation(x)
+        if self.upsample:
+            h = self.upsample(h)
+            x = self.upsample(x)
+        h = self.conv1(h)
+        
+        h = self.activation(h)
+        h = self.conv2(h)
+        
+        if self.learnable_sc:
+            x = self.conv_sc(x)
+
+
+        if self.skip_connection:
+            out = h + x
+        else:
+            out = h
+        return out
+
+def content_encoder_arch(ch =64,out_channel_multiplier = 1, input_nc = 3):
+    arch = {}
+    n=2
+    arch[80] = {'in_channels':   [input_nc] + [ch*item for item in  [1,2]],
+                                'out_channels' : [item * ch for item in [1,2,4]],
+                                'resolution': [40,20,10]}
+    arch[96] = {'in_channels':   [input_nc] + [ch*item for item in  [1,2]],
+                                'out_channels' : [item * ch for item in [1,2,4]],
+                                'resolution': [48,24,12]}
+                                
+    arch[128] = {'in_channels':   [input_nc] + [ch*item for item in  [1,2,4,8]],
+                                'out_channels' : [item * ch for item in [1,2,4,8,16]],
+                                'resolution': [64,32,16,8,4]}
+    
+    arch[256] = {'in_channels':[input_nc]+[ch*item for item in [1,2,4,8,8]],
+                                'out_channels':[item*ch for item in [1,2,4,8,8,16]],
+                                'resolution': [128,64,32,16,8,4]}
+    return arch
+
+class ContentEncoder(ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(self, G_ch=64, G_wide=True, resolution=128, 
+                 G_kernel_size=3, G_attn='64_32_16_8', n_classes=1000, 
+                 num_G_SVs=1, num_G_SV_itrs=1, G_activation=nn.ReLU(inplace=False), 
+                 SN_eps=1e-12, output_dim=1,  G_fp16=False, 
+                 G_init='N02',  G_param='SN', nf_mlp = 512, nEmbedding = 256, input_nc = 3,output_nc = 3):
+        super(ContentEncoder, self).__init__()
+
+        self.ch = G_ch
+        self.G_wide = G_wide
+        self.resolution = resolution
+        self.kernel_size = G_kernel_size
+        self.attention = G_attn
+        self.n_classes = n_classes
+        self.activation = G_activation
+        self.init = G_init
+        self.G_param = G_param
+        self.SN_eps = SN_eps
+        self.fp16 = G_fp16
+
+        if self.resolution == 96:
+            self.save_featrues = [0,1,2,3,4]
+        elif self.resolution == 80:
+            self.save_featrues = [0,1,2,3,4]
+        elif self.resolution == 128:
+            self.save_featrues = [0,1,2,3,4]
+        elif self.resolution == 256:
+            self.save_featrues = [0,1,2,3,4,5]
+        
+        self.out_channel_nultipiler = 1
+        self.arch = content_encoder_arch(self.ch, self.out_channel_nultipiler,input_nc)[resolution]
+
+        if self.G_param == 'SN':
+            self.which_conv = functools.partial(SNConv2d,
+                                                    kernel_size=3, padding=1,
+                                                    num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
+                                                    eps=self.SN_eps)
+            self.which_linear = functools.partial(SNLinear,
+                                                    num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
+                                                    eps=self.SN_eps)
+        self.blocks = []
+        for index in range(len(self.arch['out_channels'])):
+
+            self.blocks += [[DBlock(in_channels=self.arch['in_channels'][index],
+                                             out_channels=self.arch['out_channels'][index],
+                                             which_conv=self.which_conv,
+                                             wide=self.G_wide,
+                                             activation=self.activation,
+                                             preactivation=(index > 0),
+                                             downsample=nn.AvgPool2d(2))]]
+
+        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
+        self.init_weights()
+
+
+    def init_weights(self):
+        self.param_count = 0
+        for module in self.modules():
+            if (isinstance(module, nn.Conv2d)
+                    or isinstance(module, nn.Linear)
+                    or isinstance(module, nn.Embedding)):
+                if self.init == 'ortho':
+                    init.orthogonal_(module.weight)
+                elif self.init == 'N02':
+                    init.normal_(module.weight, 0, 0.02)
+                elif self.init in ['glorot', 'xavier']:
+                    init.xavier_uniform_(module.weight)
+                else:
+                    print('Init style not recognized...')
+                self.param_count += sum([p.data.nelement() for p in module.parameters()])
+        print('Param count for D''s initialized parameters: %d' % self.param_count)
+
+    def forward(self,x):
+        h = x
+        residual_features = []
+        residual_features.append(h)
+        for index, blocklist in enumerate(self.blocks):
+            for block in blocklist:
+                h = block(h)            
+            if index in self.save_featrues[:-1]:
+                residual_features.append(h)        
+        return h,residual_features

src/modules/embeddings.py

@@ -0,0 +1,84 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param embedding_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.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, channel: int, time_embed_dim: int, act_fn: str = "silu"):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(channel, time_embed_dim)
+        self.act = None
+        if act_fn == "silu":
+            self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)
+
+    def forward(self, sample):
+        sample = self.linear_1(sample)
+
+        if self.act is not None:
+            sample = self.act(sample)
+
+        sample = self.linear_2(sample)
+        return sample
+
+
+class Timesteps(nn.Module):
+    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+
+    def forward(self, timesteps):
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+        )
+        return t_emb

src/modules/resnet.py

@@ -0,0 +1,353 @@
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def upfirdn2d_native(tensor, kernel, up=1, down=1, pad=(0, 0)):
+    up_x = up_y = up
+    down_x = down_y = down
+    pad_x0 = pad_y0 = pad[0]
+    pad_x1 = pad_y1 = pad[1]
+
+    _, channel, in_h, in_w = tensor.shape
+    tensor = tensor.reshape(-1, in_h, in_w, 1)
+
+    _, in_h, in_w, minor = tensor.shape
+    kernel_h, kernel_w = kernel.shape
+
+    out = tensor.view(-1, in_h, 1, in_w, 1, minor)
+
+    # Temporary workaround for mps specific issue: https://github.com/pytorch/pytorch/issues/84535
+    if tensor.device.type == "mps":
+        out = out.to("cpu")
+    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+    out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+    out = out.to(tensor.device)  # Move back to mps if necessary
+    out = out[
+        :,
+        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+        :,
+    ]
+
+    out = out.permute(0, 3, 1, 2)
+    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
+    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+    out = F.conv2d(out, w)
+    out = out.reshape(
+        -1,
+        minor,
+        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+    )
+    out = out.permute(0, 2, 3, 1)
+    out = out[:, ::down_y, ::down_x, :]
+
+    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+
+    return out.view(-1, channel, out_h, out_w)
+
+
+def upsample_2d(hidden_states, kernel=None, factor=2, gain=1):
+    r"""Upsample2D a batch of 2D images with the given filter.
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given
+    filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified
+    `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its shape is
+    a: multiple of the upsampling factor.
+
+    Args:
+        hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        kernel: FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to nearest-neighbor upsampling.
+        factor: Integer upsampling factor (default: 2).
+        gain: Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        output: Tensor of the shape `[N, C, H * factor, W * factor]`
+    """
+    assert isinstance(factor, int) and factor >= 1
+    if kernel is None:
+        kernel = [1] * factor
+
+    kernel = torch.tensor(kernel, dtype=torch.float32)
+    if kernel.ndim == 1:
+        kernel = torch.outer(kernel, kernel)
+    kernel /= torch.sum(kernel)
+
+    kernel = kernel * (gain * (factor**2))
+    pad_value = kernel.shape[0] - factor
+    output = upfirdn2d_native(
+        hidden_states,
+        kernel.to(device=hidden_states.device),
+        up=factor,
+        pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),
+    )
+    return output
+
+
+def downsample_2d(hidden_states, kernel=None, factor=2, gain=1):
+    r"""Downsample2D a batch of 2D images with the given filter.
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and downsamples each image with the
+    given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the
+    specified `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its
+    shape is a multiple of the downsampling factor.
+
+    Args:
+        hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        kernel: FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to average pooling.
+        factor: Integer downsampling factor (default: 2).
+        gain: Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        output: Tensor of the shape `[N, C, H // factor, W // factor]`
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    if kernel is None:
+        kernel = [1] * factor
+
+    kernel = torch.tensor(kernel, dtype=torch.float32)
+    if kernel.ndim == 1:
+        kernel = torch.outer(kernel, kernel)
+    kernel /= torch.sum(kernel)
+
+    kernel = kernel * gain
+    pad_value = kernel.shape[0] - factor
+    output = upfirdn2d_native(
+        hidden_states, kernel.to(device=hidden_states.device), down=factor, pad=((pad_value + 1) // 2, pad_value // 2)
+    )
+    return output
+
+
+class Mish(torch.nn.Module):
+    def forward(self, hidden_states):
+        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))
+
+
+class Downsample2D(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    Parameters:
+        channels: channels in the inputs and outputs.
+        use_conv: a bool determining if a convolution is applied.
+        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=False, out_channels=None, padding=1, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = 2
+        self.name = name
+
+        if use_conv:
+            conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
+        else:
+            assert self.channels == self.out_channels
+            conv = nn.AvgPool2d(kernel_size=stride, stride=stride)
+
+        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states):
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv and self.padding == 0:
+            pad = (0, 1, 0, 1)
+            hidden_states = F.pad(hidden_states, pad, mode="constant", value=0)
+
+        assert hidden_states.shape[1] == self.channels
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlock2D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout=0.0,
+        temb_channels=512,
+        groups=32,
+        groups_out=None,
+        pre_norm=True,
+        eps=1e-6,
+        non_linearity="swish",
+        time_embedding_norm="default",
+        kernel=None,
+        output_scale_factor=1.0,
+        use_in_shortcut=None,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.time_embedding_norm = time_embedding_norm
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+
+        if groups_out is None:
+            groups_out = groups
+
+        self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if temb_channels is not None:
+            self.time_emb_proj = torch.nn.Linear(temb_channels, out_channels)
+        else:
+            self.time_emb_proj = None
+
+        self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+
+        self.upsample = self.downsample = None
+        if self.up:
+            if kernel == "fir":
+                fir_kernel = (1, 3, 3, 1)
+                self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)
+            elif kernel == "sde_vp":
+                self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+            else:
+                self.upsample = Upsample2D(in_channels, use_conv=False)
+        elif self.down:
+            if kernel == "fir":
+                fir_kernel = (1, 3, 3, 1)
+                self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)
+            elif kernel == "sde_vp":
+                self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)
+            else:
+                self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")
+
+        self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)   # hidden_states: torch.Size([1, 128, 64, 64])
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:  # when crossattn, both upsample and downsample is None
+            input_tensor = self.upsample(input_tensor)
+            hidden_states = self.upsample(hidden_states)
+        elif self.downsample is not None:
+            input_tensor = self.downsample(input_tensor)
+            hidden_states = self.downsample(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
+            hidden_states = hidden_states + temb   # just add together
+
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+class Upsample2D(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    Parameters:
+        channels: channels in the inputs and outputs.
+        use_conv: a bool determining if a convolution is applied.
+        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=False, use_conv_transpose=False, out_channels=None, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        conv = None
+        if use_conv_transpose:
+            conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1)
+        elif use_conv:
+            conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1)
+
+        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(self, hidden_states, output_size=None):
+        assert hidden_states.shape[1] == self.channels
+
+        if self.use_conv_transpose:
+            return self.conv(hidden_states)
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
+        # https://github.com/pytorch/pytorch/issues/86679
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if output_size is None:
+            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+        else:
+            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+        if self.use_conv:
+            if self.name == "conv":
+                hidden_states = self.conv(hidden_states)
+            else:
+                hidden_states = self.Conv2d_0(hidden_states)
+
+        return hidden_states

src/modules/style_encoder.py

@@ -0,0 +1,442 @@
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+from diffusers import ModelMixin
+from diffusers.configuration_utils import (ConfigMixin, 
+                                           register_to_config)
+
+
+def proj(x, y):
+    return torch.mm(y, x.t()) * y / torch.mm(y, y.t())
+
+
+def gram_schmidt(x, ys):
+  for y in ys:
+      x = x - proj(x, y)
+  return x
+
+
+def power_iteration(W, u_, update=True, eps=1e-12):
+    us, vs, svs = [], [], []
+    for i, u in enumerate(u_):
+        with torch.no_grad():
+            v = torch.matmul(u, W)
+            v = F.normalize(gram_schmidt(v, vs), eps=eps)
+            vs += [v]
+            u = torch.matmul(v, W.t())
+            u = F.normalize(gram_schmidt(u, us), eps=eps)
+            us += [u]
+            if update:
+                u_[i][:] = u
+        svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]
+    return svs, us, vs
+
+
+class LinearBlock(nn.Module):
+    def __init__(
+        self, 
+        in_dim, 
+        out_dim, 
+        norm='none', 
+        act='relu', 
+        use_sn=False
+    ):
+        super(LinearBlock, self).__init__()
+        use_bias = True
+        self.fc = nn.Linear(in_dim, out_dim, bias=use_bias)
+        if use_sn:
+            self.fc = nn.utils.spectral_norm(self.fc)
+
+        # initialize normalization
+        norm_dim = out_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm1d(norm_dim)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm1d(norm_dim)
+        elif norm == 'none':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if act == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif act == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif act == 'tanh':
+            self.activation = nn.Tanh()
+        elif act == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(act)
+
+    def forward(self, x):
+        out = self.fc(x)
+        if self.norm:
+            out = self.norm(out)
+        if self.activation:
+            out = self.activation(out)
+        return out
+
+
+class MLP(nn.Module):
+    def __init__(
+        self, 
+        nf_in, 
+        nf_out, 
+        nf_mlp, 
+        num_blocks, 
+        norm, 
+        act, 
+        use_sn =False
+    ):
+        super(MLP,self).__init__()
+        self.model = nn.ModuleList()
+        nf = nf_mlp
+        self.model.append(LinearBlock(nf_in, nf, norm = norm, act = act, use_sn = use_sn))
+        for _ in range((num_blocks - 2)):
+            self.model.append(LinearBlock(nf, nf, norm=norm, act=act, use_sn=use_sn))
+        self.model.append(LinearBlock(nf, nf_out, norm='none', act ='none', use_sn = use_sn))
+        self.model = nn.Sequential(*self.model)
+    
+    def forward(self, x):
+        return self.model(x.view(x.size(0), -1))
+
+
+class SN(object):
+    def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):
+        self.num_itrs = num_itrs
+        self.num_svs = num_svs
+        self.transpose = transpose
+        self.eps = eps
+        for i in range(self.num_svs):
+            self.register_buffer('u%d' % i, torch.randn(1, num_outputs))
+            self.register_buffer('sv%d' % i, torch.ones(1))
+
+    @property
+    def u(self):
+        return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]
+
+    @property
+    def sv(self):
+        return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]
+
+    def W_(self):
+        W_mat = self.weight.view(self.weight.size(0), -1)
+        if self.transpose:
+            W_mat = W_mat.t()
+        for _ in range(self.num_itrs):
+            svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)
+        if self.training:
+            with torch.no_grad():
+                for i, sv in enumerate(svs):
+                    self.sv[i][:] = sv
+        return self.weight / svs[0]
+
+
+class SNConv2d(nn.Conv2d, SN):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                padding=0, dilation=1, groups=1, bias=True,
+                num_svs=1, num_itrs=1, eps=1e-12):
+        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride,
+                        padding, dilation, groups, bias)
+        SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)
+    
+    def forward(self, x):
+        return F.conv2d(x, self.W_(), self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+    
+    def forward_wo_sn(self, x):
+        return F.conv2d(x, self.weight, self.bias, self.stride,
+                    self.padding, self.dilation, self.groups)
+
+
+class SNLinear(nn.Linear, SN):
+    def __init__(self, in_features, out_features, bias=True,
+                num_svs=1, num_itrs=1, eps=1e-12):
+        nn.Linear.__init__(self, in_features, out_features, bias)
+        SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)
+    
+    def forward(self, x):
+        return F.linear(x, self.W_(), self.bias)
+
+
+class DBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,
+                preactivation=False, activation=None, downsample=None,):
+        super(DBlock, self).__init__()
+    
+        self.in_channels, self.out_channels = in_channels, out_channels
+        
+        self.hidden_channels = self.out_channels if wide else self.in_channels
+        self.which_conv = which_conv
+        self.preactivation = preactivation
+        self.activation = activation
+        self.downsample = downsample
+
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)
+        self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)
+        self.learnable_sc = True if (in_channels != out_channels) or downsample else False
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                          kernel_size=1, padding=0)
+    def shortcut(self, x):
+        if self.preactivation:
+            if self.learnable_sc:
+                x = self.conv_sc(x)
+            if self.downsample:
+                x = self.downsample(x)
+        else:
+            if self.downsample:
+                x = self.downsample(x)
+            if self.learnable_sc:
+                x = self.conv_sc(x)
+        return x
+
+    def forward(self, x):
+    
+        if self.preactivation:
+            h = F.relu(x)
+        else:
+            h = x
+        h = self.conv1(h)
+        h = self.conv2(self.activation(h))
+        if self.downsample:
+            h = self.downsample(h)
+
+        return h + self.shortcut(x)
+
+
+class GBlock(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 which_conv=nn.Conv2d,which_bn= nn.BatchNorm2d, activation=None,
+                 upsample=None):
+        super(GBlock, self).__init__()
+
+        self.in_channels, self.out_channels = in_channels, out_channels
+        self.which_conv,self.which_bn =which_conv, which_bn
+        self.activation = activation
+        self.upsample = upsample
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.out_channels)
+        self.conv2 = self.which_conv(self.out_channels, self.out_channels)
+        self.learnable_sc = in_channels != out_channels or upsample
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                           kernel_size=1, padding=0)
+        # Batchnorm layers
+        self.bn1 = self.which_bn(in_channels)
+        self.bn2 = self.which_bn(out_channels)
+        # upsample layers
+        self.upsample = upsample
+
+    
+    def forward(self, x):
+        h = self.activation(self.bn1(x))
+        if self.upsample:
+            h = self.upsample(h)
+            x = self.upsample(x)
+        h = self.conv1(h)
+        h = self.activation(self.bn2(h))
+        h = self.conv2(h)
+        if self.learnable_sc:
+            x = self.conv_sc(x)
+        return h + x
+
+
+class GBlock2(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                which_conv=nn.Conv2d, activation=None,
+                upsample=None, skip_connection = True):
+        super(GBlock2, self).__init__()
+
+        self.in_channels, self.out_channels = in_channels, out_channels
+        self.which_conv = which_conv
+        self.activation = activation
+        self.upsample = upsample
+
+        # Conv layers
+        self.conv1 = self.which_conv(self.in_channels, self.out_channels)
+        self.conv2 = self.which_conv(self.out_channels, self.out_channels)
+        self.learnable_sc = in_channels != out_channels or upsample
+        if self.learnable_sc:
+            self.conv_sc = self.which_conv(in_channels, out_channels,
+                                          kernel_size=1, padding=0)
+        # upsample layers
+        self.upsample = upsample
+        self.skip_connection = skip_connection
+
+    def forward(self, x):
+        h = self.activation(x)
+        if self.upsample:
+            h = self.upsample(h)
+            x = self.upsample(x)
+        h = self.conv1(h)
+        
+        h = self.activation(h)
+        h = self.conv2(h)
+        
+        if self.learnable_sc:
+            x = self.conv_sc(x)
+        if self.skip_connection:
+            out = h + x
+        else:
+            out = h
+        return out
+
+
+def style_encoder_textedit_addskip_arch(ch =64,out_channel_multiplier = 1, input_nc = 3):
+    arch = {}
+    n=2
+    arch[96] = {'in_channels':   [input_nc] + [ch*item for item in  [1,2,4,8]],
+                                'out_channels' : [item * ch for item in [1,2,4,8,16]],
+                                'resolution': [48,24,12,6,3]}
+
+    arch[128] = {'in_channels':   [input_nc] + [ch*item for item in  [1,2,4,8]],
+                                'out_channels' : [item * ch for item in [1,2,4,8,16]],
+                                'resolution': [64,32,16,8,4]}
+    
+    arch[256] = {'in_channels':[input_nc]+[ch*item for item in [1,2,4,8,8]],
+                                'out_channels':[item*ch for item in [1,2,4,8,8,16]],
+                                'resolution': [128,64,32,16,8,4]}
+    return arch
+
+
+class StyleEncoder(ModelMixin, ConfigMixin):
+    """
+    This class is to encode the style image to image embedding.
+    Downsample scale is 32.
+    For example:
+        Input: Shape[Batch, 3, 128, 128]
+        Output: Shape[Batch, 255, 4, 4]
+    """
+    @register_to_config
+    def __init__(
+        self, 
+        G_ch=64, 
+        G_wide=True, 
+        resolution=128,
+        G_kernel_size=3, 
+        G_attn='64_32_16_8', 
+        n_classes=1000,
+        num_G_SVs=1, 
+        num_G_SV_itrs=1, 
+        G_activation=nn.ReLU(inplace=False),
+        SN_eps=1e-12, 
+        output_dim=1,  
+        G_fp16=False,
+        G_init='N02',  
+        G_param='SN', 
+        nf_mlp = 512, 
+        nEmbedding = 256, 
+        input_nc = 3,
+        output_nc = 3
+    ):
+        super(StyleEncoder, self).__init__()
+
+        self.ch = G_ch
+        self.G_wide = G_wide
+        self.resolution = resolution
+        self.kernel_size = G_kernel_size
+        self.attention = G_attn
+        self.n_classes = n_classes
+        self.activation = G_activation
+        self.init = G_init
+        self.G_param = G_param
+        self.SN_eps = SN_eps
+        self.fp16 = G_fp16
+
+        if self.resolution == 96:
+            self.save_featrues = [0,1,2,3,4]
+        if self.resolution == 128:
+            self.save_featrues = [0,1,2,3,4]
+        elif self.resolution == 256:
+            self.save_featrues = [0,1,2,3,4,5]
+        
+        self.out_channel_nultipiler = 1
+        self.arch = style_encoder_textedit_addskip_arch(
+          self.ch, 
+          self.out_channel_nultipiler,
+          input_nc
+        )[resolution]
+
+        if self.G_param == 'SN':
+            self.which_conv = functools.partial(
+              SNConv2d,
+              kernel_size=3, padding=1,
+              num_svs=num_G_SVs, 
+              num_itrs=num_G_SV_itrs,
+              eps=self.SN_eps
+            )
+            self.which_linear = functools.partial(
+              SNLinear,
+              num_svs=num_G_SVs, 
+              num_itrs=num_G_SV_itrs,
+              eps=self.SN_eps
+            )
+        self.blocks = []
+        for index in range(len(self.arch['out_channels'])):
+
+            self.blocks += [[DBlock(
+              in_channels=self.arch['in_channels'][index],
+              out_channels=self.arch['out_channels'][index],
+              which_conv=self.which_conv,
+              wide=self.G_wide,
+              activation=self.activation,
+              preactivation=(index > 0),
+              downsample=nn.AvgPool2d(2)
+            )]]
+
+        self.blocks = nn.ModuleList([
+          nn.ModuleList(block) for block in self.blocks
+        ])
+        last_layer = nn.Sequential(
+          nn.InstanceNorm2d(self.arch['out_channels'][-1]),
+          self.activation,
+          nn.Conv2d(
+            self.arch['out_channels'][-1],
+            self.arch['out_channels'][-1],
+            kernel_size=1,
+            stride=1
+          )
+        )
+        self.blocks.append(last_layer)
+        self.init_weights()
+
+    def init_weights(self):
+        self.param_count = 0
+        for module in self.modules():
+            if (isinstance(module, nn.Conv2d)
+                    or isinstance(module, nn.Linear)
+                    or isinstance(module, nn.Embedding)):
+                if self.init == 'ortho':
+                    init.orthogonal_(module.weight)
+                elif self.init == 'N02':
+                    init.normal_(module.weight, 0, 0.02)
+                elif self.init in ['glorot', 'xavier']:
+                    init.xavier_uniform_(module.weight)
+                else:
+                    print('Init style not recognized...')
+                self.param_count += sum([p.data.nelement() for p in module.parameters()])
+        print('Param count for D''s initialized parameters: %d' % self.param_count)
+
+    def forward(self,x):        
+        h = x
+        residual_features = []
+        residual_features.append(h)
+        for index, blocklist in enumerate(self.blocks):
+            for block in blocklist:
+                h = block(h)            
+            if index in self.save_featrues[:-1]:
+                residual_features.append(h)        
+        h = self.blocks[-1](h)
+        style_emd = h        
+        h = F.adaptive_avg_pool2d(h,(1,1))
+        h = h.view(h.size(0),-1)
+        
+        return style_emd,h,residual_features

src/modules/unet.py

@@ -0,0 +1,299 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers import ModelMixin
+from diffusers.configuration_utils import (ConfigMixin, 
+                                           register_to_config)
+from diffusers.utils import BaseOutput, logging
+
+from .embeddings import TimestepEmbedding, Timesteps
+from .unet_blocks import (DownBlock2D,
+                          UNetMidMCABlock2D,
+                          UpBlock2D,
+                          get_down_block,
+                          get_up_block)
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class UNetOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class UNet(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = None,
+        up_block_types: Tuple[str] = None,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: int = 1,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1280,
+        attention_head_dim: int = 8,
+        channel_attn: bool = False,
+        content_encoder_downsample_size: int = 4,
+        content_start_channel: int = 16,
+        reduction: int = 32,
+    ):
+        super().__init__()
+
+        self.content_encoder_downsample_size = content_encoder_downsample_size
+
+        self.sample_size = sample_size
+        time_embed_dim = block_out_channels[0] * 4
+
+        # input
+        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))
+
+        # time
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            if i != 0: 
+                content_channel = content_start_channel * (2 ** (i-1))
+            else:
+                content_channel = 0
+
+            print("Load the down block ", down_block_type)
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim,
+                downsample_padding=downsample_padding,
+                content_channel=content_channel,
+                reduction=reduction,
+                channel_attn=channel_attn,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidMCABlock2D(
+            in_channels=block_out_channels[-1],
+            temb_channels=time_embed_dim,
+            channel_attn=channel_attn,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            resnet_time_scale_shift="default",
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attention_head_dim,
+            resnet_groups=norm_num_groups,
+            content_channel=content_start_channel*(2**(content_encoder_downsample_size - 1)),
+            reduction=reduction,
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            content_channel = content_start_channel * (2 ** (content_encoder_downsample_size - i - 1))
+
+            print("Load the up block ", up_block_type)
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=layers_per_block + 1,  # larger 1 than the down block
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim,
+                upblock_index=i,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.config.attention_head_dim % slice_size != 0:
+            raise ValueError(
+                f"Make sure slice_size {slice_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.config.attention_head_dim}"
+            )
+        if slice_size is not None and slice_size > self.config.attention_head_dim:
+            raise ValueError(
+                f"Chunk_size {slice_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.config.attention_head_dim}"
+            )
+
+        for block in self.down_blocks:
+            if hasattr(block, "attentions") and block.attentions is not None:
+                block.set_attention_slice(slice_size)
+
+        self.mid_block.set_attention_slice(slice_size)
+
+        for block in self.up_blocks:
+            if hasattr(block, "attentions") and block.attentions is not None:
+                block.set_attention_slice(slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (DownBlock2D, UpBlock2D)):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        content_encoder_downsample_size: int = 4,
+        return_dict: bool = False,
+    ) -> Union[UNetOutput, Tuple]:
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # 1. time
+        timesteps = timestep   # only one time
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)
+        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+        emb = self.time_embedding(t_emb)  # projection
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for index, downsample_block in enumerate(self.down_blocks):
+            if (hasattr(downsample_block, "attentions") and downsample_block.attentions is not None) or hasattr(downsample_block, "content_attentions"):
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    index=index,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)   
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample, 
+                emb, 
+                index=content_encoder_downsample_size,
+                encoder_hidden_states=encoder_hidden_states
+            )
+
+        # 5. up
+        offset_out_sum = 0
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if (hasattr(upsample_block, "attentions") and upsample_block.attentions is not None) or hasattr(upsample_block, "content_attentions"):
+                sample, offset_out = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    style_structure_features=encoder_hidden_states[3],
+                    encoder_hidden_states=encoder_hidden_states[2],
+                )
+                offset_out_sum += offset_out
+            else:
+                sample = upsample_block(
+                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+                )
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample, offset_out_sum)
+
+        return UNetOutput(sample=sample)

src/modules/unet_blocks.py

@@ -0,0 +1,661 @@
+import torch
+from torch import nn
+from torchvision.ops import DeformConv2d
+
+from .attention import (SpatialTransformer, 
+                        OffsetRefStrucInter, 
+                        ChannelAttnBlock)
+from .resnet import (Downsample2D, 
+                     ResnetBlock2D, 
+                     Upsample2D)
+
+
+def get_down_block(
+    down_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    temb_channels,
+    add_downsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    downsample_padding=None,
+    channel_attn=False,
+    content_channel=32,
+    reduction=32):
+
+    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
+    if down_block_type == "DownBlock2D":
+        return DownBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding)
+    elif down_block_type == "MCADownBlock2D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
+        return MCADownBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            channel_attn=channel_attn,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            content_channel=content_channel,
+            reduction=reduction)
+    else:
+        raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block(
+    up_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    prev_output_channel,
+    temb_channels,
+    add_upsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    upblock_index,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    structure_feature_begin=64):
+
+    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    if up_block_type == "UpBlock2D":
+        return UpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups)
+    elif up_block_type == "StyleRSIUpBlock2D":
+        return StyleRSIUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            structure_feature_begin=structure_feature_begin,
+            upblock_index=upblock_index)
+    else:
+        raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidMCABlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        channel_attn: bool = False,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        attention_type="default",
+        output_scale_factor=1.0,
+        cross_attention_dim=1280,
+        content_channel=256,
+        reduction=32,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        content_attentions = []
+        style_attentions = []
+
+        for _ in range(num_layers):
+            content_attentions.append(
+                ChannelAttnBlock(
+                    in_channels=in_channels + content_channel,
+                    out_channels=in_channels,
+                    non_linearity=resnet_act_fn,
+                    channel_attn=channel_attn,
+                    reduction=reduction,
+                )
+            )
+            style_attentions.append(
+                SpatialTransformer(
+                    in_channels,
+                    attn_num_head_channels,
+                    in_channels // attn_num_head_channels,
+                    depth=1,
+                    context_dim=cross_attention_dim,
+                    num_groups=resnet_groups,
+                )
+            )
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.content_attentions = nn.ModuleList(content_attentions)
+        self.style_attentions = nn.ModuleList(style_attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(
+        self, 
+        hidden_states, 
+        temb=None, 
+        encoder_hidden_states=None,
+        index=None,
+    ):
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for content_attn, style_attn, resnet in zip(self.content_attentions, self.style_attentions, self.resnets[1:]):
+            
+            # content
+            current_content_feature = encoder_hidden_states[1][index]
+            hidden_states = content_attn(hidden_states, current_content_feature)
+            
+            # t_embed
+            hidden_states = resnet(hidden_states, temb)
+
+            # style
+            current_style_feature = encoder_hidden_states[0]
+            batch_size, channel, height, width = current_style_feature.shape
+            current_style_feature = current_style_feature.permute(0, 2, 3, 1).reshape(batch_size, height*width, channel)
+            hidden_states = style_attn(hidden_states, context=current_style_feature)
+
+        return hidden_states
+
+
+class MCADownBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        channel_attn: bool = False,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        attention_type="default",
+        output_scale_factor=1.0,
+        downsample_padding=1,
+        add_downsample=True,
+        content_channel=16,
+        reduction=32,
+    ):
+        super().__init__()
+        content_attentions = []
+        resnets = []
+        style_attentions = []
+
+        self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            content_attentions.append(
+                ChannelAttnBlock(
+                    in_channels=in_channels+content_channel,
+                    out_channels=in_channels,
+                    groups=resnet_groups,
+                    non_linearity=resnet_act_fn,
+                    channel_attn=channel_attn,
+                    reduction=reduction,
+                )
+            )
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            print("The style_attention cross attention dim in Down Block {} layer is {}".format(i+1, cross_attention_dim))
+            style_attentions.append(
+                SpatialTransformer(
+                    out_channels,
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    depth=1,
+                    context_dim=cross_attention_dim,
+                    num_groups=resnet_groups,
+                )
+            )
+        self.content_attentions = nn.ModuleList(content_attentions)
+        self.style_attentions = nn.ModuleList(style_attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        if num_layers == 1:
+            in_channels = out_channels
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self, 
+        hidden_states, 
+        index,
+        temb=None, 
+        encoder_hidden_states=None
+    ):
+        output_states = ()
+
+        for content_attn, resnet, style_attn in zip(self.content_attentions, self.resnets, self.style_attentions):
+            
+            # content
+            current_content_feature = encoder_hidden_states[1][index]
+            hidden_states = content_attn(hidden_states, current_content_feature)
+            
+            # t_embed
+            hidden_states = resnet(hidden_states, temb)
+
+            # style
+            current_style_feature = encoder_hidden_states[0]
+            batch_size, channel, height, width = current_style_feature.shape
+            current_style_feature = current_style_feature.permute(0, 2, 3, 1).reshape(batch_size, height*width, channel)
+            hidden_states = style_attn(hidden_states, context=current_style_feature)
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class DownBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_downsample=True,
+        downsample_padding=1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if num_layers == 1:
+            in_channels = out_channels
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, temb=None):
+        output_states = ()
+
+        for resnet in self.resnets:
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class StyleRSIUpBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        attention_type="default",
+        output_scale_factor=1.0,
+        downsample_padding=1,
+        structure_feature_begin=64, 
+        upblock_index=1,
+        add_upsample=True,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        sc_interpreter_offsets = []
+        dcn_deforms = []
+
+        self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
+        self.upblock_index = upblock_index
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+            
+            sc_interpreter_offsets.append(
+                OffsetRefStrucInter(
+                    res_in_channels=res_skip_channels,
+                    style_feat_in_channels=int(structure_feature_begin * 2 / upblock_index),
+                    n_heads=attn_num_head_channels,
+                    num_groups=resnet_groups,
+                )
+            )
+            dcn_deforms.append(
+                DeformConv2d(
+                    in_channels=res_skip_channels,
+                    out_channels=res_skip_channels,
+                    kernel_size=(3, 3),
+                    stride=1,
+                    padding=1,
+                    dilation=1,
+                )
+            )
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            attentions.append(
+                SpatialTransformer(
+                    out_channels,
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    depth=1,
+                    context_dim=cross_attention_dim,
+                    num_groups=resnet_groups,
+                )
+            )
+        self.sc_interpreter_offsets = nn.ModuleList(sc_interpreter_offsets)
+        self.dcn_deforms = nn.ModuleList(dcn_deforms)
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.num_layers = num_layers
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.attn_num_head_channels % slice_size != 0:
+            raise ValueError(
+                f"Make sure slice_size {slice_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+        if slice_size is not None and slice_size > self.attn_num_head_channels:
+            raise ValueError(
+                f"Chunk_size {slice_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+
+        for attn in self.attentions:
+            attn._set_attention_slice(slice_size)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        res_hidden_states_tuple,
+        style_structure_features,
+        temb=None,
+        encoder_hidden_states=None,
+        upsample_size=None,
+    ):
+        total_offset = 0
+
+        style_content_feat = style_structure_features[-self.upblock_index-2]
+
+        for i, (sc_inter_offset, dcn_deform, resnet, attn) in \
+            enumerate(zip(self.sc_interpreter_offsets, self.dcn_deforms, self.resnets, self.attentions)):
+            # pop res hidden states 
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            
+            # Skip Style Content Interpreter by DCN
+            offset = sc_inter_offset(res_hidden_states, style_content_feat)
+            offset = offset.contiguous()
+            # offset sum
+            offset_sum = torch.mean(torch.abs(offset))
+            total_offset += offset_sum
+
+            res_hidden_states = res_hidden_states.contiguous()
+            res_hidden_states = dcn_deform(res_hidden_states, offset)
+            # concat as input
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn), hidden_states, encoder_hidden_states
+                )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = attn(hidden_states, context=encoder_hidden_states)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        offset_out = total_offset / self.num_layers    
+
+        return hidden_states, offset_out
+
+
+class UpBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_output_channel: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_upsample=True,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
+        for resnet in self.resnets:
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states

ttf/KaiXinSongA.ttf

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

ttf/KaiXinSongB.ttf

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

utils.py

@@ -0,0 +1,123 @@
+import os
+import cv2
+import yaml
+import copy
+import pygame
+import numpy as np
+from PIL import Image
+from fontTools.ttLib import TTFont
+
+import torch
+import torchvision.transforms as transforms
+
+def save_args_to_yaml(args, output_file):
+    # Convert args namespace to a dictionary
+    args_dict = vars(args)
+
+    # Write the dictionary to a YAML file
+    with open(output_file, 'w') as yaml_file:
+        yaml.dump(args_dict, yaml_file, default_flow_style=False)
+
+
+def save_single_image(save_dir, image):
+
+    save_path = f"{save_dir}/out_single.png"
+    image.save(save_path)
+
+
+def save_image_with_content_style(save_dir, image, content_image_pil, content_image_path, style_image_path, resolution):
+    
+    new_image = Image.new('RGB', (resolution*3, resolution))
+    if content_image_pil is not None:
+        content_image = content_image_pil
+    else:
+        content_image = Image.open(content_image_path).convert("RGB").resize((resolution, resolution), Image.BILINEAR)
+    style_image = Image.open(style_image_path).convert("RGB").resize((resolution, resolution), Image.BILINEAR)
+
+    new_image.paste(content_image, (0, 0))
+    new_image.paste(style_image, (resolution, 0))
+    new_image.paste(image, (resolution*2, 0))
+
+    save_path = f"{save_dir}/out_with_cs.jpg"
+    new_image.save(save_path)
+
+
+def x0_from_epsilon(scheduler, noise_pred, x_t, timesteps):
+    """Return the x_0 from epsilon
+    """
+    batch_size = noise_pred.shape[0]
+    for i in range(batch_size):
+        noise_pred_i = noise_pred[i]
+        noise_pred_i = noise_pred_i[None, :]
+        t = timesteps[i]
+        x_t_i = x_t[i]
+        x_t_i = x_t_i[None, :]
+
+        pred_original_sample_i = scheduler.step(
+            model_output=noise_pred_i,
+            timestep=t,
+            sample=x_t_i,
+            # predict_epsilon=True,
+            generator=None,
+            return_dict=True,
+        ).pred_original_sample
+        if i == 0:
+            pred_original_sample = pred_original_sample_i
+        else:
+            pred_original_sample = torch.cat((pred_original_sample, pred_original_sample_i), dim=0)
+
+    return pred_original_sample
+
+
+def reNormalize_img(pred_original_sample):
+    pred_original_sample = (pred_original_sample / 2 + 0.5).clamp(0, 1)
+    
+    return pred_original_sample
+
+
+def normalize_mean_std(image):
+    transforms_norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    image = transforms_norm(image)
+
+    return image
+
+
+def is_char_in_font(font_path, char):
+    TTFont_font = TTFont(font_path)
+    cmap = TTFont_font['cmap']
+    for subtable in cmap.tables:
+        if ord(char) in subtable.cmap:
+            return True
+    return False
+
+
+def load_ttf(ttf_path, fsize=128):
+    pygame.init()
+
+    font = pygame.freetype.Font(ttf_path, size=fsize)
+    return font
+
+
+def ttf2im(font, char, fsize=128):
+    
+    try:
+        surface, _ = font.render(char)
+    except:
+        print("No glyph for char {}".format(char))
+        return
+    bg = np.full((fsize, fsize), 255)
+    imo = pygame.surfarray.pixels_alpha(surface).transpose(1, 0)
+    imo = 255 - np.array(Image.fromarray(imo))
+    im = copy.deepcopy(bg)
+    h, w = imo.shape[:2]
+    if h > fsize:
+        h, w = fsize, round(w*fsize/h)
+        imo = cv2.resize(imo, (w, h))
+    if w > fsize:
+        h, w = round(h*fsize/w), fsize
+        imo = cv2.resize(imo, (w, h))
+    x, y = round((fsize-w)/2), round((fsize-h)/2)
+    im[y:h+y, x:x+w] = imo
+    pil_im = Image.fromarray(im.astype('uint8')).convert('RGB')
+    
+    return pil_im