init comparison

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
 import torch
 
-from diffusers import AutoPipelineForInpainting, UNet2DConditionModel
+from diffusers import AutoPipelineForInpainting
 import diffusers
 from share_btn import community_icon_html, loading_icon_html, share_js
+from sdxl import gen_sdxl_simplified_interface
+from sdxl_models import SDXLUNet, SDXLVae, SDXLControlNetPreEncodedControlnetCond
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to(device)
 
+comparing_unet = SDXLUNet.load_fp16(device=device)
+comparing_vae = SDXLVae.load_fp16_fix(device=device)
+comparing_controlnet = SDXLControlNetPreEncodedControlnetCond.load("", device="cuda") # TODO - upload checkpoint
+comparing_controlnet.to(torch.float16)
+
 def read_content(file_path: str) -> str:
     """read the content of target file
     """
@@ -34,8 +41,12 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
     mask = dict["mask"].convert("RGB").resize((1024, 1024))
     
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
+    output_controlnet_vae_encoding = gen_sdxl_simplified_interface(
+        prompt=prompt, negative_prompt=negative_prompt, images=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps),
+        text_encoder_one=pipe.text_encoder, text_encoder_two=pipe.text_encoder_2, unet=comparing_unet, vae=comparing_vae, controlnet=comparing_controlnet, device=device
+    )
     
-    return output.images[0], gr.update(visible=True)
+    return output.images[0], output_controlnet_vae_encoding[0], gr.update(visible=True)
 
 
 css = '''
@@ -98,14 +109,15 @@ with image_blocks as demo:
                         
                 with gr.Column():
                     image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+                    image_out_comparing = gr.Image(label="Output", elem_id="output-img-comparing", height=400)
                     with gr.Group(elem_id="share-btn-container", visible=False) as share_btn_container:
                         community_icon = gr.HTML(community_icon_html)
                         loading_icon = gr.HTML(loading_icon_html)
                         share_button = gr.Button("Share to community", elem_id="share-btn",visible=True)
             
 
-    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container], api_name='run')
-    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container])
+    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container], api_name='run')
+    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container])
     share_button.click(None, [], [], _js=share_js)
 
     gr.Examples(

diffusion.py

@@ -0,0 +1,58 @@
+import torch
+
+default_num_train_timesteps = 1000
+
+
+@torch.no_grad()
+def make_sigmas(beta_start=0.00085, beta_end=0.012, num_train_timesteps=default_num_train_timesteps, device=None):
+    betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32, device=device) ** 2
+
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+
+    # TODO - would be nice to use a direct expression for this
+    sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+
+    return sigmas
+
+
+@torch.no_grad()
+def rk_ode_solver_diffusion_loop(eps_theta, timesteps, sigmas, x_T, rk_steps_weights):
+    x_t = x_T
+
+    for i in range(len(timesteps) - 1, -1, -1):
+        t = timesteps[i]
+
+        sigma = sigmas[i]
+
+        if i == 0:
+            eps_hat = eps_theta(x_t=x_t, t=t, sigma=sigma)
+            x_0_hat = x_t - sigma * eps_hat
+        else:
+            dt = sigmas[i - 1] - sigma
+
+            dx_by_dt = torch.zeros_like(x_t)
+            dx_by_dt_cur = torch.zeros_like(x_t)
+
+            for rk_step, rk_weight in rk_steps_weights:
+                dt_ = dt * rk_step
+                t_ = t + dt_
+                x_t_ = x_t + dx_by_dt_cur * dt_
+                eps_hat = eps_theta(x_t=x_t_, t=t_, sigma=sigma)
+                # TODO - note which specific ode this is the solution to and
+                # how input scaling does/doesn't effect the solution
+                dx_by_dt_cur = (x_t_ - sigma * eps_hat) / sigma
+                dx_by_dt += dx_by_dt_cur * rk_weight
+
+            x_t_minus_1 = x_t + dx_by_dt * dt
+
+            x_t = x_t_minus_1
+
+    return x_0_hat
+
+
+euler_ode_solver_diffusion_loop = lambda *args, **kwargs: rk_ode_solver_diffusion_loop(*args, **kwargs, rk_steps_weights=[[0, 1]])
+
+heun_ode_solver_diffusion_loop = lambda *args, **kwargs: rk_ode_solver_diffusion_loop(*args, **kwargs, rk_steps_weights=[[0, 0.5], [1, 0.5]])
+
+rk4_ode_solver_diffusion_loop = lambda *args, **kwargs: rk_ode_solver_diffusion_loop(*args, **kwargs, rk_steps_weights=[[0, 1 / 6], [1 / 2, 1 / 3], [1 / 2, 1 / 3], [1, 1 / 6]])

load_state_dict_patch.py

@@ -0,0 +1,415 @@
+import itertools
+from collections import OrderedDict
+from typing import Any, List, Mapping
+
+import torch
+from torch.nn import Module
+from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, _IncompatibleKeys
+
+# fmt: off
+
+# this patch is for adding the `assign` key to load_state_dict.
+# the code is in pytorch source for version 2.1
+
+def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                          missing_keys, unexpected_keys, error_msgs):
+    r"""Copies parameters and buffers from :attr:`state_dict` into only
+    this module, but not its descendants. This is called on every submodule
+    in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
+    module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
+    For state dicts without metadata, :attr:`local_metadata` is empty.
+    Subclasses can achieve class-specific backward compatible loading using
+    the version number at `local_metadata.get("version", None)`.
+    Additionally, :attr:`local_metadata` can also contain the key
+    `assign_to_params_buffers` that indicates whether keys should be
+    assigned their corresponding tensor in the state_dict.
+
+    .. note::
+        :attr:`state_dict` is not the same object as the input
+        :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
+        it can be modified.
+
+    Args:
+        state_dict (dict): a dict containing parameters and
+            persistent buffers.
+        prefix (str): the prefix for parameters and buffers used in this
+            module
+        local_metadata (dict): a dict containing the metadata for this module.
+            See
+        strict (bool): whether to strictly enforce that the keys in
+            :attr:`state_dict` with :attr:`prefix` match the names of
+            parameters and buffers in this module
+        missing_keys (list of str): if ``strict=True``, add missing keys to
+            this list
+        unexpected_keys (list of str): if ``strict=True``, add unexpected
+            keys to this list
+        error_msgs (list of str): error messages should be added to this
+            list, and will be reported together in
+            :meth:`~torch.nn.Module.load_state_dict`
+    """
+    for hook in self._load_state_dict_pre_hooks.values():
+        hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+    persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set}
+    local_name_params = itertools.chain(self._parameters.items(), persistent_buffers.items())
+    local_state = {k: v for k, v in local_name_params if v is not None}
+    assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)
+
+    for name, param in local_state.items():
+        key = prefix + name
+        if key in state_dict:
+            input_param = state_dict[key]
+            if not torch.overrides.is_tensor_like(input_param):
+                error_msgs.append('While copying the parameter named "{}", '
+                                  'expected torch.Tensor or Tensor-like object from checkpoint but '
+                                  'received {}'
+                                  .format(key, type(input_param)))
+                continue
+
+            # This is used to avoid copying uninitialized parameters into
+            # non-lazy modules, since they dont have the hook to do the checks
+            # in such case, it will error when accessing the .shape attribute.
+            is_param_lazy = torch.nn.parameter.is_lazy(param)
+            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
+            if not is_param_lazy and len(param.shape) == 0 and len(input_param.shape) == 1:
+                input_param = input_param[0]
+
+            if not is_param_lazy and input_param.shape != param.shape:
+                # local shape should match the one in checkpoint
+                error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, '
+                                  'the shape in current model is {}.'
+                                  .format(key, input_param.shape, param.shape))
+                continue
+            try:
+                with torch.no_grad():
+                    if assign_to_params_buffers:
+                        # Shape checks are already done above
+                        if (isinstance(param, torch.nn.Parameter) and
+                                not isinstance(input_param, torch.nn.Parameter)):
+                            setattr(self, name, torch.nn.Parameter(input_param))
+                        else:
+                            setattr(self, name, input_param)
+                    else:
+                        param.copy_(input_param)
+            except Exception as ex:
+                error_msgs.append('While copying the parameter named "{}", '
+                                  'whose dimensions in the model are {} and '
+                                  'whose dimensions in the checkpoint are {}, '
+                                  'an exception occurred : {}.'
+                                  .format(key, param.size(), input_param.size(), ex.args))
+        elif strict:
+            missing_keys.append(key)
+
+    extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+    if getattr(self.__class__, "set_extra_state", Module.set_extra_state) is not Module.set_extra_state:
+        if extra_state_key in state_dict:
+            self.set_extra_state(state_dict[extra_state_key])
+        elif strict:
+            missing_keys.append(extra_state_key)
+    elif strict and (extra_state_key in state_dict):
+        unexpected_keys.append(extra_state_key)
+
+    if strict:
+        for key in state_dict.keys():
+            if key.startswith(prefix) and key != extra_state_key:
+                input_name = key[len(prefix):]
+                input_name = input_name.split('.', 1)[0]  # get the name of param/buffer/child
+                if input_name not in self._modules and input_name not in local_state:
+                    unexpected_keys.append(key)
+
+def load_state_dict(self, state_dict: Mapping[str, Any],
+                    strict: bool = True, assign: bool = False):
+    r"""Copies parameters and buffers from :attr:`state_dict` into
+    this module and its descendants. If :attr:`strict` is ``True``, then
+    the keys of :attr:`state_dict` must exactly match the keys returned
+    by this module's :meth:`~torch.nn.Module.state_dict` function.
+
+    .. warning::
+        If :attr:`assign` is ``True`` the optimizer must be created after
+        the call to :attr:`load_state_dict`.
+
+    Args:
+        state_dict (dict): a dict containing parameters and
+            persistent buffers.
+        strict (bool, optional): whether to strictly enforce that the keys
+            in :attr:`state_dict` match the keys returned by this module's
+            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
+        assign (bool, optional): whether to assign items in the state
+            dictionary to their corresponding keys in the module instead
+            of copying them inplace into the module's current parameters and buffers.
+            When ``False``, the properties of the tensors in the current
+            module are preserved while when ``True``, the properties of the
+            Tensors in the state dict are preserved.
+            Default: ``False``
+
+    Returns:
+        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+            * **missing_keys** is a list of str containing the missing keys
+            * **unexpected_keys** is a list of str containing the unexpected keys
+
+    Note:
+        If a parameter or buffer is registered as ``None`` and its corresponding key
+        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
+        ``RuntimeError``.
+    """
+    if not isinstance(state_dict, Mapping):
+        raise TypeError("Expected state_dict to be dict-like, got {}.".format(type(state_dict)))
+
+    missing_keys: List[str] = []
+    unexpected_keys: List[str] = []
+    error_msgs: List[str] = []
+
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, '_metadata', None)
+    state_dict = OrderedDict(state_dict)
+    if metadata is not None:
+        # mypy isn't aware that "_metadata" exists in state_dict
+        state_dict._metadata = metadata  # type: ignore[attr-defined]
+
+    def load(module, local_state_dict, prefix=''):
+        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+        if assign:
+            local_metadata['assign_to_params_buffers'] = assign
+        module._load_from_state_dict(
+            local_state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+        for name, child in module._modules.items():
+            if child is not None:
+                child_prefix = prefix + name + '.'
+                child_state_dict = {k: v for k, v in local_state_dict.items() if k.startswith(child_prefix)}
+                load(child, child_state_dict, child_prefix)
+
+        # Note that the hook can modify missing_keys and unexpected_keys.
+        incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys)
+        for hook in module._load_state_dict_post_hooks.values():
+            out = hook(module, incompatible_keys)
+            assert out is None, (
+                "Hooks registered with ``register_load_state_dict_post_hook`` are not"
+                "expected to return new values, if incompatible_keys need to be modified,"
+                "it should be done inplace."
+            )
+
+    load(self, state_dict)
+    del load
+
+    if strict:
+        if len(unexpected_keys) > 0:
+            error_msgs.insert(
+                0, 'Unexpected key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+        if len(missing_keys) > 0:
+            error_msgs.insert(
+                0, 'Missing key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in missing_keys)))
+
+    if len(error_msgs) > 0:
+        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
+                           self.__class__.__name__, "\n\t".join(error_msgs)))
+    return _IncompatibleKeys(missing_keys, unexpected_keys)
+
+if [int(x) for x in torch.__version__.split('.')[0:2]] < [2, 1]:
+    Module._load_from_state_dict = _load_from_state_dict
+    Module.load_state_dict = load_state_dict
+
+# this patch is for adding the `assign` key to load_state_dict.
+# the code is in pytorch source for version 2.1
+
+def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                          missing_keys, unexpected_keys, error_msgs):
+    r"""Copies parameters and buffers from :attr:`state_dict` into only
+    this module, but not its descendants. This is called on every submodule
+    in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
+    module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
+    For state dicts without metadata, :attr:`local_metadata` is empty.
+    Subclasses can achieve class-specific backward compatible loading using
+    the version number at `local_metadata.get("version", None)`.
+    Additionally, :attr:`local_metadata` can also contain the key
+    `assign_to_params_buffers` that indicates whether keys should be
+    assigned their corresponding tensor in the state_dict.
+
+    .. note::
+        :attr:`state_dict` is not the same object as the input
+        :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
+        it can be modified.
+
+    Args:
+        state_dict (dict): a dict containing parameters and
+            persistent buffers.
+        prefix (str): the prefix for parameters and buffers used in this
+            module
+        local_metadata (dict): a dict containing the metadata for this module.
+            See
+        strict (bool): whether to strictly enforce that the keys in
+            :attr:`state_dict` with :attr:`prefix` match the names of
+            parameters and buffers in this module
+        missing_keys (list of str): if ``strict=True``, add missing keys to
+            this list
+        unexpected_keys (list of str): if ``strict=True``, add unexpected
+            keys to this list
+        error_msgs (list of str): error messages should be added to this
+            list, and will be reported together in
+            :meth:`~torch.nn.Module.load_state_dict`
+    """
+    for hook in self._load_state_dict_pre_hooks.values():
+        hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+    persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set}
+    local_name_params = itertools.chain(self._parameters.items(), persistent_buffers.items())
+    local_state = {k: v for k, v in local_name_params if v is not None}
+    assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)
+
+    for name, param in local_state.items():
+        key = prefix + name
+        if key in state_dict:
+            input_param = state_dict[key]
+            if not torch.overrides.is_tensor_like(input_param):
+                error_msgs.append('While copying the parameter named "{}", '
+                                  'expected torch.Tensor or Tensor-like object from checkpoint but '
+                                  'received {}'
+                                  .format(key, type(input_param)))
+                continue
+
+            # This is used to avoid copying uninitialized parameters into
+            # non-lazy modules, since they dont have the hook to do the checks
+            # in such case, it will error when accessing the .shape attribute.
+            is_param_lazy = torch.nn.parameter.is_lazy(param)
+            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
+            if not is_param_lazy and len(param.shape) == 0 and len(input_param.shape) == 1:
+                input_param = input_param[0]
+
+            if not is_param_lazy and input_param.shape != param.shape:
+                # local shape should match the one in checkpoint
+                error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, '
+                                  'the shape in current model is {}.'
+                                  .format(key, input_param.shape, param.shape))
+                continue
+            try:
+                with torch.no_grad():
+                    if assign_to_params_buffers:
+                        # Shape checks are already done above
+                        if (isinstance(param, torch.nn.Parameter) and
+                                not isinstance(input_param, torch.nn.Parameter)):
+                            setattr(self, name, torch.nn.Parameter(input_param))
+                        else:
+                            setattr(self, name, input_param)
+                    else:
+                        param.copy_(input_param)
+            except Exception as ex:
+                error_msgs.append('While copying the parameter named "{}", '
+                                  'whose dimensions in the model are {} and '
+                                  'whose dimensions in the checkpoint are {}, '
+                                  'an exception occurred : {}.'
+                                  .format(key, param.size(), input_param.size(), ex.args))
+        elif strict:
+            missing_keys.append(key)
+
+    extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+    if getattr(self.__class__, "set_extra_state", Module.set_extra_state) is not Module.set_extra_state:
+        if extra_state_key in state_dict:
+            self.set_extra_state(state_dict[extra_state_key])
+        elif strict:
+            missing_keys.append(extra_state_key)
+    elif strict and (extra_state_key in state_dict):
+        unexpected_keys.append(extra_state_key)
+
+    if strict:
+        for key in state_dict.keys():
+            if key.startswith(prefix) and key != extra_state_key:
+                input_name = key[len(prefix):]
+                input_name = input_name.split('.', 1)[0]  # get the name of param/buffer/child
+                if input_name not in self._modules and input_name not in local_state:
+                    unexpected_keys.append(key)
+
+def load_state_dict(self, state_dict: Mapping[str, Any],
+                    strict: bool = True, assign: bool = False):
+    r"""Copies parameters and buffers from :attr:`state_dict` into
+    this module and its descendants. If :attr:`strict` is ``True``, then
+    the keys of :attr:`state_dict` must exactly match the keys returned
+    by this module's :meth:`~torch.nn.Module.state_dict` function.
+
+    .. warning::
+        If :attr:`assign` is ``True`` the optimizer must be created after
+        the call to :attr:`load_state_dict`.
+
+    Args:
+        state_dict (dict): a dict containing parameters and
+            persistent buffers.
+        strict (bool, optional): whether to strictly enforce that the keys
+            in :attr:`state_dict` match the keys returned by this module's
+            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
+        assign (bool, optional): whether to assign items in the state
+            dictionary to their corresponding keys in the module instead
+            of copying them inplace into the module's current parameters and buffers.
+            When ``False``, the properties of the tensors in the current
+            module are preserved while when ``True``, the properties of the
+            Tensors in the state dict are preserved.
+            Default: ``False``
+
+    Returns:
+        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+            * **missing_keys** is a list of str containing the missing keys
+            * **unexpected_keys** is a list of str containing the unexpected keys
+
+    Note:
+        If a parameter or buffer is registered as ``None`` and its corresponding key
+        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
+        ``RuntimeError``.
+    """
+    if not isinstance(state_dict, Mapping):
+        raise TypeError("Expected state_dict to be dict-like, got {}.".format(type(state_dict)))
+
+    missing_keys: List[str] = []
+    unexpected_keys: List[str] = []
+    error_msgs: List[str] = []
+
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, '_metadata', None)
+    state_dict = OrderedDict(state_dict)
+    if metadata is not None:
+        # mypy isn't aware that "_metadata" exists in state_dict
+        state_dict._metadata = metadata  # type: ignore[attr-defined]
+
+    def load(module, local_state_dict, prefix=''):
+        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+        if assign:
+            local_metadata['assign_to_params_buffers'] = assign
+        module._load_from_state_dict(
+            local_state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+        for name, child in module._modules.items():
+            if child is not None:
+                child_prefix = prefix + name + '.'
+                child_state_dict = {k: v for k, v in local_state_dict.items() if k.startswith(child_prefix)}
+                load(child, child_state_dict, child_prefix)
+
+        # Note that the hook can modify missing_keys and unexpected_keys.
+        incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys)
+        for hook in module._load_state_dict_post_hooks.values():
+            out = hook(module, incompatible_keys)
+            assert out is None, (
+                "Hooks registered with ``register_load_state_dict_post_hook`` are not"
+                "expected to return new values, if incompatible_keys need to be modified,"
+                "it should be done inplace."
+            )
+
+    load(self, state_dict)
+    del load
+
+    if strict:
+        if len(unexpected_keys) > 0:
+            error_msgs.insert(
+                0, 'Unexpected key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+        if len(missing_keys) > 0:
+            error_msgs.insert(
+                0, 'Missing key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in missing_keys)))
+
+    if len(error_msgs) > 0:
+        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
+                           self.__class__.__name__, "\n\t".join(error_msgs)))
+    return _IncompatibleKeys(missing_keys, unexpected_keys)
+
+if [int(x) for x in torch.__version__.split('.')[0:2]] < [2, 1]:
+    Module._load_from_state_dict = _load_from_state_dict
+    Module.load_state_dict = load_state_dict
+
+# fmt: on

sdxl.py

@@ -0,0 +1,962 @@
+import itertools
+import os
+import random
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+
+import numpy as np
+import safetensors.torch
+import torch
+import torch.nn.functional as F
+import torchvision.transforms
+import torchvision.transforms.functional as TF
+import wandb
+import webdataset as wds
+from PIL import Image
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import default_collate
+from transformers import (CLIPTextModel, CLIPTextModelWithProjection,
+                          CLIPTokenizerFast)
+
+from diffusion import (default_num_train_timesteps,
+                       euler_ode_solver_diffusion_loop, make_sigmas)
+from sdxl_models import (SDXLAdapter, SDXLControlNet, SDXLControlNetFull,
+                         SDXLControlNetPreEncodedControlnetCond, SDXLUNet,
+                         SDXLVae)
+
+
+class SDXLTraining:
+    text_encoder_one: CLIPTextModel
+    text_encoder_two: CLIPTextModelWithProjection
+    vae: SDXLVae
+    sigmas: torch.Tensor
+    unet: SDXLUNet
+    adapter: Optional[SDXLAdapter]
+    controlnet: Optional[Union[SDXLControlNet, SDXLControlNetFull]]
+
+    train_unet: bool
+    train_unet_up_blocks: bool
+
+    mixed_precision: Optional[torch.dtype]
+    timestep_sampling: Literal["uniform", "cubic"]
+
+    validation_images_logged: bool
+    log_validation_input_images_every_time: bool
+
+    get_sdxl_conditioning_images: Callable[[Image.Image], Dict[str, Any]]
+
+    def __init__(
+        self,
+        device,
+        train_unet,
+        get_sdxl_conditioning_images,
+        train_unet_up_blocks=False,
+        unet_resume_from=None,
+        controlnet_cls=None,
+        controlnet_resume_from=None,
+        adapter_cls=None,
+        adapter_resume_from=None,
+        mixed_precision=None,
+        timestep_sampling="uniform",
+        log_validation_input_images_every_time=True,
+    ):
+        self.text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16)
+        self.text_encoder_one.to(device=device)
+        self.text_encoder_one.requires_grad_(False)
+        self.text_encoder_one.eval()
+
+        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2", variant="fp16", torch_dtype=torch.float16)
+        self.text_encoder_two.to(device=device)
+        self.text_encoder_two.requires_grad_(False)
+        self.text_encoder_two.eval()
+
+        self.vae = SDXLVae.load_fp16_fix(device=device)
+        self.vae.requires_grad_(False)
+        self.vae.eval()
+
+        self.sigmas = make_sigmas(device=device)
+
+        if train_unet:
+            if unet_resume_from is None:
+                self.unet = SDXLUNet.load_fp32(device=device)
+            else:
+                self.unet = SDXLUNet.load(unet_resume_from, device=device)
+            self.unet.requires_grad_(True)
+            self.unet.train()
+            self.unet = DDP(self.unet, device_ids=[device])
+        elif train_unet_up_blocks:
+            if unet_resume_from is None:
+                self.unet = SDXLUNet.load_fp32(device=device)
+            else:
+                self.unet = SDXLUNet.load_fp32(device=device, overrides=[unet_resume_from])
+            self.unet.requires_grad_(False)
+            self.unet.eval()
+            self.unet.up_blocks.requires_grad_(True)
+            self.unet.up_blocks.train()
+            self.unet = DDP(self.unet, device_ids=[device], find_unused_parameters=True)
+        else:
+            self.unet = SDXLUNet.load_fp16(device=device)
+            self.unet.requires_grad_(False)
+            self.unet.eval()
+
+        if controlnet_cls is not None:
+            if controlnet_resume_from is None:
+                self.controlnet = controlnet_cls.from_unet(self.unet)
+                self.controlnet.to(device)
+            else:
+                self.controlnet = controlnet_cls.load(controlnet_resume_from, device=device)
+            self.controlnet.train()
+            self.controlnet.requires_grad_(True)
+            # TODO add back
+            # controlnet.enable_gradient_checkpointing()
+            # TODO - should be able to remove find_unused_parameters. Comes from pre encoded controlnet
+            self.controlnet = DDP(self.controlnet, device_ids=[device], find_unused_parameters=True)
+        else:
+            self.controlnet = None
+
+        if adapter_cls is not None:
+            if adapter_resume_from is None:
+                self.adapter = adapter_cls()
+                self.adapter.to(device=device)
+            else:
+                self.adapter = adapter_cls.load(adapter_resume_from, device=device)
+            self.adapter.train()
+            self.adapter.requires_grad_(True)
+            self.adapter = DDP(self.adapter, device_ids=[device])
+        else:
+            self.adapter = None
+
+        self.mixed_precision = mixed_precision
+        self.timestep_sampling = timestep_sampling
+
+        self.validation_images_logged = False
+        self.log_validation_input_images_every_time = log_validation_input_images_every_time
+
+        self.get_sdxl_conditioning_images = get_sdxl_conditioning_images
+
+        self.train_unet = train_unet
+        self.train_unet_up_blocks = train_unet_up_blocks
+
+    def train_step(self, batch):
+        with torch.no_grad():
+            if isinstance(self.unet, DDP):
+                unet_dtype = self.unet.module.dtype
+                unet_device = self.unet.module.device
+            else:
+                unet_dtype = self.unet.dtype
+                unet_device = self.unet.device
+
+            micro_conditioning = batch["micro_conditioning"].to(device=unet_device)
+
+            image = batch["image"].to(self.vae.device, dtype=self.vae.dtype)
+            latents = self.vae.encode(image).to(dtype=unet_dtype)
+
+            text_input_ids_one = batch["text_input_ids_one"].to(self.text_encoder_one.device)
+            text_input_ids_two = batch["text_input_ids_two"].to(self.text_encoder_two.device)
+
+            encoder_hidden_states, pooled_encoder_hidden_states = sdxl_text_conditioning(self.text_encoder_one, self.text_encoder_two, text_input_ids_one, text_input_ids_two)
+
+            encoder_hidden_states = encoder_hidden_states.to(dtype=unet_dtype)
+            pooled_encoder_hidden_states = pooled_encoder_hidden_states.to(dtype=unet_dtype)
+
+            bsz = latents.shape[0]
+
+            if self.timestep_sampling == "uniform":
+                timesteps = torch.randint(0, default_num_train_timesteps, (bsz,), device=unet_device)
+            elif self.timestep_sampling == "cubic":
+                # Cubic sampling to sample a random timestep for each image
+                timesteps = torch.rand((bsz,), device=unet_device)
+                timesteps = (1 - timesteps**3) * default_num_train_timesteps
+                timesteps = timesteps.long()
+                timesteps = timesteps.clamp(0, default_num_train_timesteps - 1)
+            else:
+                assert False
+
+            sigmas_ = self.sigmas[timesteps].to(dtype=latents.dtype)
+
+            noise = torch.randn_like(latents)
+
+            noisy_latents = latents + noise * sigmas_
+
+            scaled_noisy_latents = noisy_latents / ((sigmas_**2 + 1) ** 0.5)
+
+            if "conditioning_image" in batch:
+                conditioning_image = batch["conditioning_image"].to(unet_device)
+
+            if self.controlnet is not None and isinstance(self.controlnet, SDXLControlNetPreEncodedControlnetCond):
+                controlnet_device = self.controlnet.module.device
+                controlnet_dtype = self.controlnet.module.dtype
+                conditioning_image = self.vae.encode(conditioning_image.to(self.vae.dtype)).to(device=controlnet_device, dtype=controlnet_dtype)
+                conditioning_image_mask = TF.resize(batch["conditioning_image_mask"], conditioning_image.shape[2:]).to(device=controlnet_device, dtype=controlnet_dtype)
+                conditioning_image = torch.concat((conditioning_image, conditioning_image_mask), dim=1)
+
+        with torch.autocast(
+            "cuda",
+            self.mixed_precision,
+            enabled=self.mixed_precision is not None,
+        ):
+            down_block_additional_residuals = None
+            mid_block_additional_residual = None
+            add_to_down_block_inputs = None
+            add_to_output = None
+
+            if self.adapter is not None:
+                down_block_additional_residuals = self.adapter(conditioning_image)
+
+            if self.controlnet is not None:
+                controlnet_out = self.controlnet(
+                    x_t=scaled_noisy_latents,
+                    t=timesteps,
+                    encoder_hidden_states=encoder_hidden_states,
+                    micro_conditioning=micro_conditioning,
+                    pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+                    controlnet_cond=conditioning_image,
+                )
+
+                down_block_additional_residuals = controlnet_out["down_block_res_samples"]
+                mid_block_additional_residual = controlnet_out["mid_block_res_sample"]
+                add_to_down_block_inputs = controlnet_out.get("add_to_down_block_inputs", None)
+                add_to_output = controlnet_out.get("add_to_output", None)
+
+            model_pred = self.unet(
+                x_t=scaled_noisy_latents,
+                t=timesteps,
+                encoder_hidden_states=encoder_hidden_states,
+                micro_conditioning=micro_conditioning,
+                pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+                down_block_additional_residuals=down_block_additional_residuals,
+                mid_block_additional_residual=mid_block_additional_residual,
+                add_to_down_block_inputs=add_to_down_block_inputs,
+                add_to_output=add_to_output,
+            ).sample
+
+            loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
+
+        return loss
+
+    @torch.no_grad()
+    def log_validation(self, step, num_validation_images: int, validation_prompts: Optional[List[str]] = None, validation_images: Optional[List[str]] = None):
+        if isinstance(self.unet, DDP):
+            unet = self.unet.module
+            unet.eval()
+            unet_set_to_eval = True
+        else:
+            unet = self.unet
+            unet_set_to_eval = False
+
+        if self.adapter is not None:
+            adapter = self.adapter.module
+            adapter.eval()
+        else:
+            adapter = None
+
+        if self.controlnet is not None:
+            controlnet = self.controlnet.module
+            controlnet.eval()
+        else:
+            controlnet = None
+
+        formatted_validation_images = None
+
+        if validation_images is not None:
+            formatted_validation_images = []
+            wandb_validation_images = []
+
+            for validation_image_path in validation_images:
+                validation_image = Image.open(validation_image_path)
+                validation_image = validation_image.convert("RGB")
+                validation_image = validation_image.resize((1024, 1024))
+
+                conditioning_images = self.get_sdxl_conditioning_images(validation_image)
+
+                conditioning_image = conditioning_images["conditioning_image"]
+
+                if self.controlnet is not None and isinstance(self.controlnet, SDXLControlNetPreEncodedControlnetCond):
+                    conditioning_image = self.vae.encode(conditioning_image[None, :, :, :].to(self.vae.device, dtype=self.vae.dtype))
+                    conditionin_mask_image = TF.resize(conditioning_images["conditioning_mask_image"], conditioning_image.shape[2:]).to(conditioning_image.dtype, conditioning_image.device)
+                    conditioning_image = torch.concat(conditioning_image, conditionin_mask_image, dim=1)
+
+                formatted_validation_images.append(conditioning_image)
+                wandb_validation_images.append(wandb.Image(conditioning_images["conditioning_image_as_pil"]))
+
+            if self.log_validation_input_images_every_time or not self.validation_images_logged:
+                wandb.log({"validation_conditioning": wandb_validation_images}, step=step)
+                self.validation_images_logged = True
+
+        generator = torch.Generator().manual_seed(0)
+
+        output_validation_images = []
+
+        for formatted_validation_image, validation_prompt in zip(formatted_validation_images, validation_prompts):
+            for _ in range(num_validation_images):
+                with torch.autocast("cuda"):
+                    x_0 = sdxl_diffusion_loop(
+                        prompts=validation_prompt,
+                        images=formatted_validation_image,
+                        unet=unet,
+                        text_encoder_one=self.text_encoder_one,
+                        text_encoder_two=self.text_encoder_two,
+                        controlnet=controlnet,
+                        adapter=adapter,
+                        sigmas=self.sigmas,
+                        generator=generator,
+                    )
+
+                    x_0 = self.vae.decode(x_0)
+                    x_0 = self.vae.output_tensor_to_pil(x_0)[0]
+
+                    output_validation_images.append(wandb.Image(x_0, caption=validation_prompt))
+
+        wandb.log({"validation": output_validation_images}, step=step)
+
+        if unet_set_to_eval:
+            unet.train()
+
+        if adapter is not None:
+            adapter.train()
+
+        if controlnet is not None:
+            controlnet.train()
+
+    def parameters(self):
+        if self.train_unet:
+            return self.unet.parameters()
+
+        if self.controlnet is not None and self.train_unet_up_blocks:
+            return itertools.chain(self.controlnet.parameters(), self.unet.up_blocks.parameters())
+
+        if self.controlnet is not None:
+            return self.controlnet.parameters()
+
+        if self.adapter is not None:
+            return self.adapter.parameters()
+
+        assert False
+
+    def save(self, save_to):
+        if self.train_unet:
+            safetensors.torch.save_file(self.unet.module.state_dict(), os.path.join(save_to, "unet.safetensors"))
+
+        if self.controlnet is not None and self.train_unet_up_blocks:
+            safetensors.torch.save_file(self.controlnet.module.state_dict(), os.path.join(save_to, "controlnet.safetensors"))
+            safetensors.torch.save_file(self.unet.module.up_blocks.state_dict(), os.path.join(save_to, "unet.safetensors"))
+
+        if self.controlnet is not None:
+            safetensors.torch.save_file(self.controlnet.module.state_dict(), os.path.join(save_to, "controlnet.safetensors"))
+
+        if self.adapter is not None:
+            safetensors.torch.save_file(self.adapter.module.state_dict(), os.path.join(save_to, "adapter.safetensors"))
+
+
+def get_sdxl_dataset(train_shards: str, shuffle_buffer_size: int, batch_size: int, proportion_empty_prompts: float, get_sdxl_conditioning_images=None):
+    dataset = (
+        wds.WebDataset(
+            train_shards,
+            resampled=True,
+            handler=wds.ignore_and_continue,
+        )
+        .shuffle(shuffle_buffer_size)
+        .decode("pil", handler=wds.ignore_and_continue)
+        .rename(
+            image="jpg;png;jpeg;webp",
+            text="text;txt;caption",
+            metadata="json",
+            handler=wds.warn_and_continue,
+        )
+        .map(lambda d: make_sample(d, proportion_empty_prompts=proportion_empty_prompts, get_sdxl_conditioning_images=get_sdxl_conditioning_images))
+        .select(lambda sample: "conditioning_image" not in sample or sample["conditioning_image"] is not None)
+    )
+
+    dataset = dataset.batched(batch_size, partial=False, collation_fn=default_collate)
+
+    return dataset
+
+
+@torch.no_grad()
+def make_sample(d, proportion_empty_prompts, get_sdxl_conditioning_images=None):
+    image = d["image"]
+    metadata = d["metadata"]
+
+    if random.random() < proportion_empty_prompts:
+        text = ""
+    else:
+        text = d["text"]
+
+    c_top, c_left, _, _ = get_random_crop_params([image.height, image.width], [1024, 1024])
+
+    original_width = int(metadata.get("original_width", 0.0))
+    original_height = int(metadata.get("original_height", 0.0))
+
+    micro_conditioning = torch.tensor([original_width, original_height, c_top, c_left, 1024, 1024])
+
+    text_input_ids_one = sdxl_tokenize_one(text)
+
+    text_input_ids_two = sdxl_tokenize_two(text)
+
+    image = image.convert("RGB")
+
+    image = TF.resize(
+        image,
+        1024,
+        interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
+    )
+
+    image = TF.crop(
+        image,
+        c_top,
+        c_left,
+        1024,
+        1024,
+    )
+
+    sample = {
+        "micro_conditioning": micro_conditioning,
+        "text_input_ids_one": text_input_ids_one,
+        "text_input_ids_two": text_input_ids_two,
+        "image": SDXLVae.input_pil_to_tensor(image),
+    }
+
+    if get_sdxl_conditioning_images is not None:
+        conditioning_images = get_sdxl_conditioning_images(image)
+
+        sample["conditioning_image"] = conditioning_images["conditioning_image"]
+
+        if conditioning_images["conditioning_image_mask"] is not None:
+            sample["conditioning_image_mask"] = conditioning_images["conditioning_image_mask"]
+
+    return sample
+
+
+def get_random_crop_params(input_size: Tuple[int, int], output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
+    h, w = input_size
+
+    th, tw = output_size
+
+    if h < th or w < tw:
+        raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
+
+    if w == tw and h == th:
+        return 0, 0, h, w
+
+    i = torch.randint(0, h - th + 1, size=(1,)).item()
+    j = torch.randint(0, w - tw + 1, size=(1,)).item()
+
+    return i, j, th, tw
+
+
+def get_sdxl_conditioning_images(image, adapter_type=None, controlnet_type=None, controlnet_variant=None, open_pose=None, conditioning_image_mask=None):
+    resolution = image.width
+
+    if adapter_type == "openpose":
+        conditioning_image = open_pose(image, detect_resolution=resolution, image_resolution=resolution, return_pil=False)
+
+        if (conditioning_image == 0).all():
+            return None, None
+
+        conditioning_image_as_pil = Image.fromarray(conditioning_image)
+
+        conditioning_image = TF.to_tensor(conditioning_image)
+
+    if controlnet_type == "canny":
+        import cv2
+
+        conditioning_image = np.array(image)
+        conditioning_image = cv2.Canny(conditioning_image, 100, 200)
+        conditioning_image = conditioning_image[:, :, None]
+        conditioning_image = np.concatenate([conditioning_image, conditioning_image, conditioning_image], axis=2)
+
+        conditioning_image_as_pil = Image.fromarray(conditioning_image)
+
+        conditioning_image = TF.to_tensor(conditioning_image)
+
+    if controlnet_type == "inpainting":
+        if conditioning_image_mask is None:
+            if random.random() <= 0.25:
+                conditioning_image_mask = np.ones((resolution, resolution), np.float32)
+            else:
+                conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
+
+            conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
+
+            conditioning_image_mask = conditioning_image_mask[None, :, :]
+
+        conditioning_image = TF.to_tensor(image)
+
+        if controlnet_variant == "pre_encoded_controlnet_cond":
+            # where mask is 1, zero out the pixels. Note that this requires mask to be concattenated
+            # with the mask so that the network knows the zeroed out pixels are from the mask and
+            # are not just zero in the original image
+            conditioning_image = conditioning_image * (conditioning_image_mask < 0.5)
+
+            conditioning_image_as_pil = TF.to_pil_image(conditioning_image)
+
+            conditioning_image = TF.normalize(conditioning_image, [0.5], [0.5])
+        else:
+            # Just zero out the pixels which will be masked
+            conditioning_image_as_pil = TF.to_pil_image(conditioning_image * (conditioning_image_mask < 0.5))
+
+            # where mask is set to 1, set to -1 "special" masked image pixel.
+            # -1 is outside of the 0-1 range that the controlnet normalized
+            # input is in.
+            conditioning_image = conditioning_image * (conditioning_image_mask < 0.5) + -1.0 * (conditioning_image_mask >= 0.5)
+
+    return dict(conditioning_image=conditioning_image, conditioning_image_mask=conditioning_image_mask, conditioning_image_as_pil=conditioning_image_as_pil)
+
+
+# TODO: would be nice to just call a function from a tokenizers https://github.com/huggingface/tokenizers
+# i.e. afaik tokenizing shouldn't require holding any state
+
+tokenizer_one = CLIPTokenizerFast.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="tokenizer")
+
+tokenizer_two = CLIPTokenizerFast.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="tokenizer_2")
+
+
+def sdxl_tokenize_one(prompts):
+    return tokenizer_one(
+        prompts,
+        padding="max_length",
+        max_length=tokenizer_one.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    ).input_ids[0]
+
+
+def sdxl_tokenize_two(prompts):
+    return tokenizer_two(
+        prompts,
+        padding="max_length",
+        max_length=tokenizer_one.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    ).input_ids[0]
+
+
+def sdxl_text_conditioning(text_encoder_one, text_encoder_two, text_input_ids_one, text_input_ids_two):
+    prompt_embeds_1 = text_encoder_one(
+        text_input_ids_one,
+        output_hidden_states=True,
+    ).hidden_states[-2]
+
+    prompt_embeds_1 = prompt_embeds_1.view(prompt_embeds_1.shape[0], prompt_embeds_1.shape[1], -1)
+
+    prompt_embeds_2 = text_encoder_two(
+        text_input_ids_two,
+        output_hidden_states=True,
+    )
+
+    pooled_encoder_hidden_states = prompt_embeds_2[0]
+
+    prompt_embeds_2 = prompt_embeds_2.hidden_states[-2]
+
+    prompt_embeds_2 = prompt_embeds_2.view(prompt_embeds_2.shape[0], prompt_embeds_2.shape[1], -1)
+
+    encoder_hidden_states = torch.cat((prompt_embeds_1, prompt_embeds_2), dim=-1)
+
+    return encoder_hidden_states, pooled_encoder_hidden_states
+
+
+def make_random_rectangle_mask(
+    height,
+    width,
+    margin=10,
+    bbox_min_size=100,
+    bbox_max_size=512,
+    min_times=1,
+    max_times=2,
+):
+    mask = np.zeros((height, width), np.float32)
+
+    bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2)
+
+    times = np.random.randint(min_times, max_times + 1)
+
+    for i in range(times):
+        box_width = np.random.randint(bbox_min_size, bbox_max_size)
+        box_height = np.random.randint(bbox_min_size, bbox_max_size)
+
+        start_x = np.random.randint(margin, width - margin - box_width + 1)
+        start_y = np.random.randint(margin, height - margin - box_height + 1)
+
+        mask[start_y : start_y + box_height, start_x : start_x + box_width] = 1
+
+    return mask
+
+
+def make_random_irregular_mask(height, width, max_angle=4, max_len=60, max_width=256, min_times=1, max_times=2):
+    import cv2
+
+    mask = np.zeros((height, width), np.float32)
+
+    times = np.random.randint(min_times, max_times + 1)
+
+    for i in range(times):
+        start_x = np.random.randint(width)
+        start_y = np.random.randint(height)
+
+        for j in range(1 + np.random.randint(5)):
+            angle = 0.01 + np.random.randint(max_angle)
+
+            if i % 2 == 0:
+                angle = 2 * 3.1415926 - angle
+
+            length = 10 + np.random.randint(max_len)
+
+            brush_w = 5 + np.random.randint(max_width)
+
+            end_x = np.clip((start_x + length * np.sin(angle)).astype(np.int32), 0, width)
+            end_y = np.clip((start_y + length * np.cos(angle)).astype(np.int32), 0, height)
+
+            choice = random.randint(0, 2)
+
+            if choice == 0:
+                cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0, brush_w)
+            elif choice == 1:
+                cv2.circle(mask, (start_x, start_y), radius=brush_w, color=1.0, thickness=-1)
+            elif choice == 2:
+                radius = brush_w // 2
+                mask[
+                    start_y - radius : start_y + radius,
+                    start_x - radius : start_x + radius,
+                ] = 1
+            else:
+                assert False
+
+            start_x, start_y = end_x, end_y
+
+    return mask
+
+
+def make_outpainting_mask(height, width, probs=[0.5, 0.5, 0.5, 0.5]):
+    mask = np.zeros((height, width), np.float32)
+    at_least_one_mask_applied = False
+
+    coords = [
+        [(0, 0), (1, get_padding(height))],
+        [(0, 0), (get_padding(width), 1)],
+        [(0, 1 - get_padding(height)), (1, 1)],
+        [(1 - get_padding(width), 0), (1, 1)],
+    ]
+
+    for pp, coord in zip(probs, coords):
+        if np.random.random() < pp:
+            at_least_one_mask_applied = True
+            mask = apply_padding(mask=mask, coord=coord)
+
+    if not at_least_one_mask_applied:
+        idx = np.random.choice(range(len(coords)), p=np.array(probs) / sum(probs))
+        mask = apply_padding(mask=mask, coord=coords[idx])
+
+    return mask
+
+
+def get_padding(size, min_padding_percent=0.04, max_padding_percent=0.5):
+    n1 = int(min_padding_percent * size)
+    n2 = int(max_padding_percent * size)
+    return np.random.randint(n1, n2) / size
+
+
+def apply_padding(mask, coord):
+    height, width = mask.shape
+
+    mask[
+        int(coord[0][0] * height) : int(coord[1][0] * height),
+        int(coord[0][1] * width) : int(coord[1][1] * width),
+    ] = 1
+
+    return mask
+
+
+@torch.no_grad()
+def sdxl_diffusion_loop(
+    prompts,
+    unet,
+    text_encoder_one,
+    text_encoder_two,
+    images=None,
+    controlnet=None,
+    adapter=None,
+    sigmas=None,
+    timesteps=None,
+    x_T=None,
+    micro_conditioning=None,
+    guidance_scale=5.0,
+    generator=None,
+    negative_prompts=None,
+    diffusion_loop=euler_ode_solver_diffusion_loop,
+):
+    if negative_prompts is None:
+        negative_prompts = [""] * len(prompts)
+
+    prompts += negative_prompts
+
+    encoder_hidden_states, pooled_encoder_hidden_states = sdxl_text_conditioning(
+        text_encoder_one,
+        text_encoder_two,
+        sdxl_tokenize_one(prompts).to(text_encoder_one.device),
+        sdxl_tokenize_two(prompts).to(text_encoder_two.device),
+    )
+
+    if x_T is None:
+        x_T = torch.randn((1, 4, 1024 // 8, 1024 // 8), dtype=torch.float32, device=unet.device, generator=generator)
+        x_T = x_T * ((sigmas.max() ** 2 + 1) ** 0.5)
+
+    if sigmas is None:
+        sigmas = make_sigmas(device=unet.device)
+
+    if timesteps is None:
+        timesteps = torch.linspace(0, sigmas.numel(), 50, dtype=torch.long, device=unet.device)
+
+    if micro_conditioning is None:
+        micro_conditioning = torch.tensor([1024, 1024, 0, 0, 1024, 1024], dtype=torch.long, device=unet.device)
+
+    if adapter is not None:
+        down_block_additional_residuals = adapter(images)
+    else:
+        down_block_additional_residuals = None
+
+    if controlnet is not None:
+        controlnet_cond = images
+    else:
+        controlnet_cond = None
+
+    eps_theta = lambda x_t, t, sigma: sdxl_eps_theta(
+        x_t=x_t,
+        t=t,
+        sigma=sigma,
+        unet=unet,
+        encoder_hidden_states=encoder_hidden_states,
+        pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+        micro_conditioning=micro_conditioning,
+        guidance_scale=guidance_scale,
+        controlnet=controlnet,
+        controlnet_cond=controlnet_cond,
+        down_block_additional_residuals=down_block_additional_residuals,
+    )
+
+    x_0 = diffusion_loop(eps_theta=eps_theta, timesteps=timesteps, sigmas=sigmas, x_T=x_T)
+
+    return x_0
+
+
+@torch.no_grad()
+def sdxl_eps_theta(
+    x_t,
+    t,
+    sigma,
+    unet,
+    encoder_hidden_states,
+    pooled_encoder_hidden_states,
+    micro_conditioning,
+    guidance_scale,
+    controlnet=None,
+    controlnet_cond=None,
+    down_block_additional_residuals=None,
+):
+    # TODO - how does this not effect the ode we are solving
+    scaled_x_t = x_t / ((sigma**2 + 1) ** 0.5)
+
+    if guidance_scale > 1.0:
+        scaled_x_t = torch.concat([scaled_x_t, scaled_x_t])
+
+    if controlnet is not None:
+        controlnet_out = controlnet(
+            x_t=scaled_x_t,
+            t=t,
+            encoder_hidden_states=encoder_hidden_states,
+            micro_conditioning=micro_conditioning,
+            pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+            controlnet_cond=controlnet_cond,
+        )
+
+        down_block_additional_residuals = controlnet_out["down_block_res_samples"]
+        mid_block_additional_residual = controlnet_out["mid_block_res_sample"]
+        add_to_down_block_inputs = controlnet_out.get("add_to_down_block_inputs", None)
+        add_to_output = controlnet_out.get("add_to_output", None)
+    else:
+        mid_block_additional_residual = None
+        add_to_down_block_inputs = None
+        add_to_output = None
+
+    eps_hat = unet(
+        x_t=scaled_x_t,
+        t=t,
+        encoder_hidden_states=encoder_hidden_states,
+        micro_conditioning=micro_conditioning,
+        pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+        down_block_additional_residuals=down_block_additional_residuals,
+        mid_block_additional_residual=mid_block_additional_residual,
+        add_to_down_block_inputs=add_to_down_block_inputs,
+        add_to_output=add_to_output,
+    )
+
+    if guidance_scale > 1.0:
+        eps_hat_uncond, eps_hat = eps_hat.chunk(2)
+
+        eps_hat = eps_hat_uncond + guidance_scale * (eps_hat - eps_hat_uncond)
+
+    return eps_hat
+
+known_negative_prompt = "text, watermark, low-quality, signature, moiré pattern, downsampling, aliasing, distorted, blurry, glossy, blur, jpeg artifacts, compression artifacts, poorly drawn, low-resolution, bad, distortion, twisted, excessive, exaggerated pose, exaggerated limbs, grainy, symmetrical, duplicate, error, pattern, beginner, pixelated, fake, hyper, glitch, overexposed, high-contrast, bad-contrast"
+
+def gen_sdxl_simplified_interface(
+    prompt:str,
+    negative_prompt: Optional[str] = None,
+    controlnet_checkpoint: Optional[str]=None,
+    controlnet: Optional[Literal["SDXLControlNet", "SDXLContolNetFull", "SDXLControlNetPreEncodedControlnetCond"]]=None,
+    adapter_checkpoint: Optional[str]=None,
+    num_inference_steps=50,
+    images=None,
+    masks=None,
+    apply_conditioning: Optional[Literal["canny"]]=None,
+    num_images: int=1,
+    device: Optional[str]=None,
+    text_encoder_one=None,
+    text_encoder_two=None,
+    unet=None,
+    vae=None,
+):
+    if device is None:
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+
+    if text_encoder_one is None:
+        text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16)
+        text_encoder_one.to(device=device)
+
+    if text_encoder_two is None:
+        text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2", variant="fp16", torch_dtype=torch.float16)
+        text_encoder_two.to(device=device)
+
+    if vae is None:
+        vae = SDXLVae.load_fp16_fix(device=device)
+
+    if unet is None:
+        unet = SDXLUNet.load_fp16(device=device)
+
+    if isinstance(controlnet, str) and controlnet_checkpoint is not None:
+        if controlnet == "SDXLControlNet":
+            controlnet = SDXLControlNet.load(controlnet_checkpoint, device=device, dtype=torch.float16)
+        elif controlnet == "SDXLControlNetFull":
+            controlnet = SDXLControlNetFull.load(controlnet_checkpoint, device=device, dtype=torch.float16)
+        elif controlnet == "SDXLControlNetPreEncodedControlnetCond":
+            controlnet = SDXLControlNetPreEncodedControlnetCond.load(controlnet_checkpoint, device=device, dtype=torch.float16)
+        else:
+            assert False
+
+    if adapter_checkpoint is not None:
+        adapter = SDXLAdapter.load(adapter_checkpoint, device=device, dtype=torch.float16)
+    else:
+        adapter = None
+
+    sigmas = make_sigmas()
+
+    timesteps = torch.linspace(0, sigmas.numel(), num_inference_steps, dtype=torch.long, device=unet.device)
+
+    if images is not None:
+        if not isinstance(images, list):
+            images = [images]
+
+        if masks is not None and not isinstance(masks, list):
+            masks = [masks]
+
+        images_ = []
+
+        for image_idx, image in enumerate(images):
+            if isinstance(image, str):
+                image = Image.open(image)
+                image = image.convert("RGB")
+                image = image.resize((1024, 1024))
+            elif isinstance(image, Image.Image):
+                ...
+            else:
+                assert False
+
+            if apply_conditioning == "canny":
+                import cv2
+
+                image = np.array(image)
+                image = cv2.Canny(image, 100, 200)
+                image = image[:, :, None]
+                controlnet_image = np.concatenate([controlnet_image, controlnet_image, controlnet_image], axis=2)
+
+            image = TF.to_tensor(image)
+
+            if masks is not None:
+                mask = masks[image_idx]
+                if isinstance(mask, str):
+                    mask = Image.open(mask)
+                    mask = mask.convert("L")
+                    mask = mask.resize((1024, 1024))
+                elif isinstance(mask, Image.Image):
+                    ...
+                else:
+                    assert False
+                mask = TF.to_tensor(mask)
+
+                if controlnet == "SDXLControlNetPreEncodedControlnetCond":
+                    image = image * (mask < 0.5)
+                    image = TF.normalized(image, [0.5], [0.5])
+                    image = vae.encode(image)
+                    mask = TF.resize(mask, (1024 // 8, 1024 // 8))
+                    image = torch.concat((image, mask))
+                else:
+                    image = image * (mask < 0.5) + -1.0 * (mask >= 0.5)
+
+            images_.append(image)
+
+        images_ = torch.concat(images_)
+    else:
+        images_ = None
+
+    x_0 = sdxl_diffusion_loop(
+        prompts=[prompt] * num_images,
+        negative_prompts=[negative_prompt] * num_images,
+        unet=unet,
+        text_encoder_one=text_encoder_one,
+        text_encoder_two=text_encoder_two,
+        sigmas=sigmas,
+        timesteps=timesteps,
+        controlnet=controlnet,
+        adapter=adapter,
+        images=images_,
+    )
+
+    x_0 = vae.decode(x_0)
+    x_0 = vae.output_tensor_to_pil(x_0)
+
+    return x_0
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    args = ArgumentParser()
+    args.add_argument("--prompt", required=True, type=str)
+    args.add_argument("--num_images", required=True, type=int, default=1)
+    args.add_argument("--num_inference_steps", required=False, type=int, default=50)
+    args.add_argument("--image", required=False, type=str, default=None)
+    args.add_argument("--mask", required=False, type=str, default=None)
+    args.add_argument("--controlnet_checkpoint", required=False, type=str, default=None)
+    args.add_argument("--controlnet", required=False, choices=["SDXLControlNet", "SDXLControlNetFull", "SDXLControNetPreEncodedControlnetCond"], default=None)
+    args.add_argument("--adapter_checkpoint", required=False, type=str, default=None)
+    args.add_argument("--apply_conditioning", choices=["canny"], required=False, default=None)
+    args.add_argument("--device", required=False, default=None)
+    args = args.parse_args()
+
+    images = gen_sdxl_simplified_interface(
+        prompt=args.prompt,
+        num_images=args.num_images,
+        num_inference_steps=args.num_inference_steps,
+        images=[args.image],
+        masks=[args.mask],
+        controlnet_checkpoint=args.controlnet_checkpoint,
+        controlnet=args.controlnet,
+        adapter_checkpoint=args.adapter_checkpoint,
+        apply_conditioning=args.apply_conditioning,
+        device=args.device,
+        negative_prompt=known_negative_prompt,
+    )
+
+    for i, image in enumerate(images):
+        image.save(f"out_{i}.png")

sdxl_models.py

@@ -0,0 +1,1375 @@
+import math
+import os
+from typing import List, Optional
+
+import safetensors.torch
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+import xformers
+from PIL import Image
+from torch import nn
+
+
+class ModelUtils:
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @classmethod
+    def load(cls, load_from: str, device, overrides: Optional[List[str]] = None):
+        import load_state_dict_patch
+
+        load_from = [load_from]
+
+        load_from += overrides
+
+        state_dict = {}
+
+        for load_from_ in load_from:
+            if os.path.isdir(load_from_):
+                load_from_ = os.path.join(load_from_, "diffusion_pytorch_model.safetensors")
+
+            state_dict.update(safetensors.torch.load_file(load_from_, device=device))
+
+        with torch.device("meta"):
+            model = cls()
+
+        model.load_state_dict(state_dict, assign=True)
+
+        return model
+
+
+vae_scaling_factor = 0.13025
+
+
+class SDXLVae(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        self.encoder = nn.ModuleDict(dict(
+            # 3 -> 128
+            conv_in=nn.Conv2d(3, 128, kernel_size=3, padding=1),
+
+            down_blocks=nn.ModuleList([
+                # 128 -> 128
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(128, 128, eps=1e-6), ResnetBlock2D(128, 128, eps=1e-6)]),
+                    downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1)))]),
+                )),
+                # 128 -> 256
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(128, 256, eps=1e-6), ResnetBlock2D(256, 256, eps=1e-6)]),
+                    downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)))]),
+                )),
+                # 256 -> 512
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(256, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
+                    downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)))]),
+                )),
+                # 512 -> 512
+                nn.ModuleDict(dict(resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]))),
+            ]),
+
+            # 512 -> 512
+            mid_block=nn.ModuleDict(dict(
+                attentions=nn.ModuleList([Attention(512, 512, qkv_bias=True)]),
+                resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
+            )),
+
+            # 512 -> 8
+            conv_norm_out=nn.GroupNorm(32, 512, eps=1e-06),
+            conv_act=nn.SiLU(),
+            conv_out=nn.Conv2d(512, 8, kernel_size=3, padding=1)
+        ))
+
+        # 8 -> 8
+        self.quant_conv = nn.Conv2d(8, 8, kernel_size=1)
+
+        # 8 -> 4 from sampling mean and std
+
+        # 4 -> 4
+        self.post_quant_conv = nn.Conv2d(4, 4, 1)
+
+        self.decoder = nn.ModuleDict(dict(
+            # 4 -> 512
+            conv_in=nn.Conv2d(4, 512, kernel_size=3, padding=1),
+
+            # 512 -> 512
+            mid_block=nn.ModuleDict(dict(
+                attentions=nn.ModuleList([Attention(512, 512, qkv_bias=True)]),
+                resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
+            )),
+
+            up_blocks=nn.ModuleList([
+                # 512 -> 512
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
+                    upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)))]),
+                )),
+
+                # 512 -> 512
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
+                    upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)))]),
+                )),
+
+                # 512 -> 256
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(512, 256, eps=1e-6), ResnetBlock2D(256, 256, eps=1e-6), ResnetBlock2D(256, 256, eps=1e-6)]),
+                    upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)))]),
+                )),
+
+                # 256 -> 128
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList([ResnetBlock2D(256, 128, eps=1e-6), ResnetBlock2D(128, 128, eps=1e-6), ResnetBlock2D(128, 128, eps=1e-6)]),
+                )),
+            ]),
+
+            # 128 -> 3
+            conv_norm_out=nn.GroupNorm(32, 128, eps=1e-06),
+            conv_act=nn.SiLU(),
+            conv_out=nn.Conv2d(128, 3, kernel_size=3, padding=1)
+        ))
+
+        # fmt: on
+
+    def encode(self, x, generator=None):
+        h = x
+
+        h = self.encoder["conv_in"](h)
+
+        for down_block in self.encoder["down_blocks"]:
+            for resnet in down_block["resnets"]:
+                h = resnet(h)
+
+            if "downsamplers" in down_block:
+                h = down_block["downsamplers"][0]["conv"](h)
+
+        h = self.encoder["mid_block"]["resnets"][0](h)
+        h = self.encoder["mid_block"]["attentions"][0](h)
+        h = self.encoder["mid_block"]["resnets"][1](h)
+
+        h = self.encoder["conv_norm_out"](h)
+        h = self.encoder["conv_act"](h)
+        h = self.encoder["conv_out"](h)
+
+        mean, logvar = self.quant_conv(h).chunk(2, dim=1)
+
+        logvar = torch.clamp(logvar, -30.0, 20.0)
+
+        std = torch.exp(0.5 * logvar)
+
+        z = mean + torch.randn(mean.shape, device=mean.device, dtype=mean.dtype, generator=generator) * std
+
+        z = z * vae_scaling_factor
+
+        return z
+
+    def decode(self, z):
+        z = z / vae_scaling_factor
+
+        h = z
+
+        h = self.post_quant_conv(h)
+
+        h = self.decoder["mid_block"]["resnets"][0](h)
+        h = self.decoder["mid_block"]["attentions"][0](h)
+        h = self.decoder["mid_block"]["resnets"][1](h)
+
+        for up_block in self.encoder["up_blocks"]:
+            for resnet in up_block["resnets"]:
+                h = resnet(h)
+
+            if "upsamplers" in up_block:
+                h = up_block["upsamplers"][0]["conv"](h)
+
+        h = self.decoder["conv_norm_out"](h)
+        h = self.decoder["conv_act"](h)
+        h = self.decoder["conv_out"](h)
+
+        x_pred = h
+
+        return x_pred
+
+    @classmethod
+    def input_pil_to_tensor(self, x):
+        x = TF.to_tensor(x)
+        x = TF.normalize(x, [0.5], [0.5])
+        if x.ndim == 3:
+            x = x[None, :, :, :]
+        return x
+
+    @classmethod
+    def output_tensor_to_pil(self, x_pred):
+        x_pred = ((x_pred * 0.5 + 0.5).clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 1)
+
+        x_pred = x_pred.permute(0, 2, 3, 1).cpu().numpy()
+
+        x_pred = [Image.fromarray(x) for x in x_pred]
+
+        return x_pred
+
+    @classmethod
+    def load_fp32(cls, device=None, overrides=None):
+        return cls.load("./weights/sdxl_vae.safetensors", device=device, overrides=overrides)
+
+    @classmethod
+    def load_fp16(cls, device=None, overrides=None):
+        return cls.load("./weights/sdxl_vae.fp16.safetensors", device=device, overrides=overrides)
+
+    @classmethod
+    def load_fp16_fix(cls, device=None, overrides=None):
+        return cls.load("./weights/sdxl_vae_fp16_fix.safetensors", device=device, overrides=overrides)
+
+
+class SDXLUNet(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        encoder_hidden_states_dim = 2048
+
+        # timesteps embedding:
+
+        time_sinusoidal_embedding_dim = 320
+        time_embedding_dim = 1280
+
+        self.get_sinusoidal_timestep_embedding = lambda timesteps: get_sinusoidal_embedding(timesteps, time_sinusoidal_embedding_dim)
+
+        self.time_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(time_sinusoidal_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # image size and crop coordinates conditioning embedding (i.e. micro conditioning):
+
+        num_micro_conditioning_values = 6
+        micro_conditioning_embedding_dim = 256
+        additional_embedding_encoder_dim = 1280
+        self.get_sinusoidal_micro_conditioning_embedding = lambda micro_conditioning: get_sinusoidal_embedding(micro_conditioning, micro_conditioning_embedding_dim)
+
+        self.add_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(additional_embedding_encoder_dim + num_micro_conditioning_values * micro_conditioning_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # actual unet blocks:
+
+        self.conv_in = nn.Conv2d(4, 320, kernel_size=3, padding=1)
+
+        self.down_blocks = nn.ModuleList([
+            # 320 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 320 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 640, time_embedding_dim),
+                    ResnetBlock2D(640, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 640 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+            )),
+        ])
+
+        self.mid_block = nn.ModuleDict(dict(
+            resnets=nn.ModuleList([
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+            ]),
+            attentions=nn.ModuleList([TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10)]),
+        ))
+
+        self.up_blocks = nn.ModuleList([
+            # 1280 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(1280 + 1280, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280 + 1280, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280 + 640, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+                upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(1280, 1280, kernel_size=3, padding=1)))]),
+            )),
+            # 1280 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(1280 + 640, 640, time_embedding_dim),
+                    ResnetBlock2D(640 + 640, 640, time_embedding_dim),
+                    ResnetBlock2D(640 + 320, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, padding=1)))]),
+            )),
+            # 640 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640 + 320, 320, time_embedding_dim),
+                    ResnetBlock2D(320 + 320, 320, time_embedding_dim),
+                    ResnetBlock2D(320 + 320, 320, time_embedding_dim),
+                ]),
+            ))
+        ])
+
+        self.conv_norm_out = nn.GroupNorm(32, 320)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(320, 4, kernel_size=3, padding=1)
+
+        # fmt: on
+
+    def forward(
+        self,
+        x_t,
+        t,
+        encoder_hidden_states,
+        micro_conditioning,
+        pooled_encoder_hidden_states,
+        down_block_additional_residuals: Optional[List[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        add_to_down_block_inputs: Optional[List[torch.Tensor]] = None,
+        add_to_output: Optional[torch.Tensor] = None,
+    ):
+        hidden_state = x_t
+
+        t = self.get_sinusoidal_timestep_embedding(t)
+        t = t.to(dtype=hidden_state.dtype)
+        t = self.time_embedding["linear_1"](t)
+        t = self.time_embedding["act"](t)
+        t = self.time_embedding["linear_2"](t)
+
+        additional_conditioning = self.get_sinusoidal_micro_conditioning_embedding(micro_conditioning)
+        additional_conditioning = additional_conditioning.to(dtype=hidden_state.dtype)
+        additional_conditioning = additional_conditioning.flatten(1)
+        additional_conditioning = torch.concat([pooled_encoder_hidden_states, additional_conditioning], dim=-1)
+        additional_conditioning = self.add_embedding["linear_1"](additional_conditioning)
+        additional_conditioning = self.add_embedding["act"](additional_conditioning)
+        additional_conditioning = self.add_embedding["linear_2"](additional_conditioning)
+
+        t = t + additional_conditioning
+
+        hidden_state = self.conv_in(hidden_state)
+
+        residuals = [hidden_state]
+
+        for down_block in self.down_blocks:
+            for i, resnet in enumerate(down_block["resnets"]):
+                if add_to_down_block_inputs is not None:
+                    hidden_state = hidden_state + add_to_down_block_inputs.pop(0)
+
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in down_block:
+                    hidden_state = down_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+                residuals.append(hidden_state)
+
+            if "downsamplers" in down_block:
+                if add_to_down_block_inputs is not None:
+                    hidden_state = hidden_state + add_to_down_block_inputs.pop(0)
+
+                hidden_state = down_block["downsamplers"][0]["conv"](hidden_state)
+
+                residuals.append(hidden_state)
+
+        hidden_state = self.mid_block["resnets"][0](hidden_state, t)
+        hidden_state = self.mid_block["attentions"][0](hidden_state, encoder_hidden_states)
+        hidden_state = self.mid_block["resnets"][1](hidden_state, t)
+
+        if mid_block_additional_residual is not None:
+            hidden_state = hidden_state + mid_block_additional_residual
+
+        for up_block in self.up_blocks:
+            for i, resnet in enumerate(up_block["resnets"]):
+                residual = residuals.pop()
+
+                if down_block_additional_residuals is not None:
+                    residual = residual + down_block_additional_residuals.pop()
+
+                hidden_state = torch.concat([hidden_state, residual], dim=1)
+
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in up_block:
+                    hidden_state = up_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+            if "upsamplers" in up_block:
+                hidden_state = F.interpolate(hidden_state, scale_factor=2.0, mode="nearest")
+                hidden_state = up_block["upsamplers"][0]["conv"](hidden_state)
+
+        hidden_state = self.conv_norm_out(hidden_state)
+        hidden_state = self.conv_act(hidden_state)
+        hidden_state = self.conv_out(hidden_state)
+
+        if add_to_output is not None:
+            hidden_state = hidden_state + add_to_output
+
+        eps_hat = hidden_state
+
+        return eps_hat
+
+    @classmethod
+    def load_fp32(cls, device=None, overrides=None):
+        return cls.load("./weights/sdxl_unet.safetensors", device=device, overrides=overrides)
+
+    @classmethod
+    def load_fp16(cls, device=None, overrides=None):
+        return cls.load("./weights/sdxl_unet.fp16.safetensors", device=device, overrides=overrides)
+
+
+class SDXLControlNet(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        encoder_hidden_states_dim = 2048
+
+        # timesteps embedding:
+
+        time_sinusoidal_embedding_dim = 320
+        time_embedding_dim = 1280
+
+        self.get_sinusoidal_timestep_embedding = lambda timesteps: get_sinusoidal_embedding(timesteps, time_sinusoidal_embedding_dim)
+
+        self.time_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(time_sinusoidal_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # image size and crop coordinates conditioning embedding (i.e. micro conditioning):
+
+        num_micro_conditioning_values = 6
+        micro_conditioning_embedding_dim = 256
+        additional_embedding_encoder_dim = 1280
+        self.get_sinusoidal_micro_conditioning_embedding = lambda micro_conditioning: get_sinusoidal_embedding(micro_conditioning, micro_conditioning_embedding_dim)
+
+        self.add_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(additional_embedding_encoder_dim + num_micro_conditioning_values * micro_conditioning_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # controlnet cond embedding:
+        self.controlnet_cond_embedding = nn.ModuleDict(dict(
+            conv_in=nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            blocks=nn.ModuleList([
+                # 16 -> 32
+                nn.Conv2d(16, 16, kernel_size=3, padding=1),
+                nn.Conv2d(16, 32, kernel_size=3, padding=1, stride=2),
+                # 32 -> 96
+                nn.Conv2d(32, 32, kernel_size=3, padding=1),
+                nn.Conv2d(32, 96, kernel_size=3, padding=1, stride=2),
+                # 96 -> 256
+                nn.Conv2d(96, 96, kernel_size=3, padding=1),
+                nn.Conv2d(96, 256, kernel_size=3, padding=1, stride=2),
+            ]),
+            conv_out=zero_module(nn.Conv2d(256, 320, kernel_size=3, padding=1)),
+        ))
+
+        # actual unet blocks:
+
+        self.conv_in = nn.Conv2d(4, 320, kernel_size=3, padding=1)
+
+        self.down_blocks = nn.ModuleList([
+            # 320 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 320 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 640, time_embedding_dim),
+                    ResnetBlock2D(640, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 640 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+            )),
+        ])
+
+        self.controlnet_down_blocks = nn.ModuleList([
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+        ])
+
+        self.mid_block = nn.ModuleDict(dict(
+            resnets=nn.ModuleList([
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+            ]),
+            attentions=nn.ModuleList([TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10)]),
+        ))
+
+        self.controlnet_mid_block = zero_module(nn.Conv2d(1280, 1280, kernel_size=1))
+
+        # fmt: on
+
+    def forward(
+        self,
+        x_t,
+        t,
+        encoder_hidden_states,
+        micro_conditioning,
+        pooled_encoder_hidden_states,
+        controlnet_cond,
+    ):
+        hidden_state = x_t
+
+        t = self.get_sinusoidal_timestep_embedding(t)
+        t = t.to(dtype=hidden_state.dtype)
+        t = self.time_embedding["linear_1"](t)
+        t = self.time_embedding["act"](t)
+        t = self.time_embedding["linear_2"](t)
+
+        additional_conditioning = self.get_sinusoidal_micro_conditioning_embedding(micro_conditioning)
+        additional_conditioning = additional_conditioning.to(dtype=hidden_state.dtype)
+        additional_conditioning = additional_conditioning.flatten(1)
+        additional_conditioning = torch.concat([pooled_encoder_hidden_states, additional_conditioning], dim=-1)
+        additional_conditioning = self.add_embedding["linear_1"](additional_conditioning)
+        additional_conditioning = self.add_embedding["act"](additional_conditioning)
+        additional_conditioning = self.add_embedding["linear_2"](additional_conditioning)
+
+        t = t + additional_conditioning
+
+        controlnet_cond = self.controlnet_cond_embedding["conv_in"](controlnet_cond)
+        controlnet_cond = F.silu(controlnet_cond)
+
+        for block in self.controlnet_cond_embedding["blocks"]:
+            controlnet_cond = F.silu(block(controlnet_cond))
+
+        controlnet_cond = self.controlnet_cond_embedding["conv_out"](controlnet_cond)
+
+        hidden_state = self.conv_in(hidden_state)
+
+        hidden_state = hidden_state + controlnet_cond
+
+        down_block_res_sample = self.controlnet_down_blocks[0](hidden_state)
+        down_block_res_samples = [down_block_res_sample]
+
+        for down_block in self.down_blocks:
+            for i, resnet in enumerate(down_block["resnets"]):
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in down_block:
+                    hidden_state = down_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+                down_block_res_sample = self.controlnet_down_blocks[len(down_block_res_samples)](hidden_state)
+                down_block_res_samples.append(down_block_res_sample)
+
+            if "downsamplers" in down_block:
+                hidden_state = down_block["downsamplers"][0]["conv"](hidden_state)
+
+                down_block_res_sample = self.controlnet_down_blocks[len(down_block_res_samples)](hidden_state)
+                down_block_res_samples.append(down_block_res_sample)
+
+        hidden_state = self.mid_block["resnets"][0](hidden_state, t)
+        hidden_state = self.mid_block["attentions"][0](hidden_state, encoder_hidden_states)
+        hidden_state = self.mid_block["resnets"][1](hidden_state, t)
+
+        mid_block_res_sample = self.controlnet_mid_block(hidden_state)
+
+        return dict(
+            down_block_res_samples=down_block_res_samples,
+            mid_block_res_sample=mid_block_res_sample,
+        )
+
+    @classmethod
+    def from_unet(cls, unet):
+        controlnet = cls()
+
+        controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+        controlnet.add_embedding.load_state_dict(unet.add_embedding.state_dict())
+
+        controlnet.conv_in.load_state_dict(unet.conv_in.state_dict())
+
+        controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())
+        controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())
+
+        return controlnet
+
+
+class SDXLControlNetPreEncodedControlnetCond(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        encoder_hidden_states_dim = 2048
+
+        # timesteps embedding:
+
+        time_sinusoidal_embedding_dim = 320
+        time_embedding_dim = 1280
+
+        self.get_sinusoidal_timestep_embedding = lambda timesteps: get_sinusoidal_embedding(timesteps, time_sinusoidal_embedding_dim)
+
+        self.time_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(time_sinusoidal_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # image size and crop coordinates conditioning embedding (i.e. micro conditioning):
+
+        num_micro_conditioning_values = 6
+        micro_conditioning_embedding_dim = 256
+        additional_embedding_encoder_dim = 1280
+        self.get_sinusoidal_micro_conditioning_embedding = lambda micro_conditioning: get_sinusoidal_embedding(micro_conditioning, micro_conditioning_embedding_dim)
+
+        self.add_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(additional_embedding_encoder_dim + num_micro_conditioning_values * micro_conditioning_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # actual unet blocks:
+
+        # unet latents: 4 + 
+        # control image latents: 4 + 
+        # controlnet_mask: 1
+        # = 9 channels
+        self.conv_in = nn.Conv2d(9, 320, kernel_size=3, padding=1)
+
+        self.down_blocks = nn.ModuleList([
+            # 320 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 320 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 640, time_embedding_dim),
+                    ResnetBlock2D(640, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 640 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+            )),
+        ])
+
+        self.controlnet_down_blocks = nn.ModuleList([
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+        ])
+
+        self.mid_block = nn.ModuleDict(dict(
+            resnets=nn.ModuleList([
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+            ]),
+            attentions=nn.ModuleList([TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10)]),
+        ))
+
+        self.controlnet_mid_block = zero_module(nn.Conv2d(1280, 1280, kernel_size=1))
+
+        # fmt: on
+
+    def forward(
+        self,
+        x_t,
+        t,
+        encoder_hidden_states,
+        micro_conditioning,
+        pooled_encoder_hidden_states,
+        controlnet_cond,
+    ):
+        hidden_state = x_t
+
+        t = self.get_sinusoidal_timestep_embedding(t)
+        t = t.to(dtype=hidden_state.dtype)
+        t = self.time_embedding["linear_1"](t)
+        t = self.time_embedding["act"](t)
+        t = self.time_embedding["linear_2"](t)
+
+        additional_conditioning = self.get_sinusoidal_micro_conditioning_embedding(micro_conditioning)
+        additional_conditioning = additional_conditioning.to(dtype=hidden_state.dtype)
+        additional_conditioning = additional_conditioning.flatten(1)
+        additional_conditioning = torch.concat([pooled_encoder_hidden_states, additional_conditioning], dim=-1)
+        additional_conditioning = self.add_embedding["linear_1"](additional_conditioning)
+        additional_conditioning = self.add_embedding["act"](additional_conditioning)
+        additional_conditioning = self.add_embedding["linear_2"](additional_conditioning)
+
+        t = t + additional_conditioning
+
+        hidden_state = torch.concat((hidden_state, controlnet_cond), dim=1)
+
+        hidden_state = self.conv_in(hidden_state)
+
+        down_block_res_sample = self.controlnet_down_blocks[0](hidden_state)
+        down_block_res_samples = [down_block_res_sample]
+
+        for down_block in self.down_blocks:
+            for i, resnet in enumerate(down_block["resnets"]):
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in down_block:
+                    hidden_state = down_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+                down_block_res_sample = self.controlnet_down_blocks[len(down_block_res_samples)](hidden_state)
+                down_block_res_samples.append(down_block_res_sample)
+
+            if "downsamplers" in down_block:
+                hidden_state = down_block["downsamplers"][0]["conv"](hidden_state)
+
+                down_block_res_sample = self.controlnet_down_blocks[len(down_block_res_samples)](hidden_state)
+                down_block_res_samples.append(down_block_res_sample)
+
+        hidden_state = self.mid_block["resnets"][0](hidden_state, t)
+        hidden_state = self.mid_block["attentions"][0](hidden_state, encoder_hidden_states)
+        hidden_state = self.mid_block["resnets"][1](hidden_state, t)
+
+        mid_block_res_sample = self.controlnet_mid_block(hidden_state)
+
+        return dict(
+            down_block_res_samples=down_block_res_samples,
+            mid_block_res_sample=mid_block_res_sample,
+        )
+
+    @classmethod
+    def from_unet(cls, unet):
+        controlnet = cls()
+
+        controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+        controlnet.add_embedding.load_state_dict(unet.add_embedding.state_dict())
+
+        conv_in_weight = unet.conv_in.state_dict()["weight"]
+        padding = torch.zeros((320, 5, 3, 3), device=conv_in_weight.device, dtype=conv_in_weight.dtype)
+        conv_in_weight = torch.concat((conv_in_weight, padding), dim=1)
+
+        conv_in_bias = unet.conv_in.state_dict()["bias"]
+
+        controlnet.conv_in.load_state_dict({"weight": conv_in_weight, "bias": conv_in_bias})
+
+        controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())
+        controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())
+
+        return controlnet
+
+
+class SDXLControlNetFull(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        encoder_hidden_states_dim = 2048
+
+        # timesteps embedding:
+
+        time_sinusoidal_embedding_dim = 320
+        time_embedding_dim = 1280
+
+        self.get_sinusoidal_timestep_embedding = lambda timesteps: get_sinusoidal_embedding(timesteps, time_sinusoidal_embedding_dim)
+
+        self.time_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(time_sinusoidal_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # image size and crop coordinates conditioning embedding (i.e. micro conditioning):
+
+        num_micro_conditioning_values = 6
+        micro_conditioning_embedding_dim = 256
+        additional_embedding_encoder_dim = 1280
+        self.get_sinusoidal_micro_conditioning_embedding = lambda micro_conditioning: get_sinusoidal_embedding(micro_conditioning, micro_conditioning_embedding_dim)
+
+        self.add_embedding = nn.ModuleDict(dict(
+            linear_1=nn.Linear(additional_embedding_encoder_dim + num_micro_conditioning_values * micro_conditioning_embedding_dim, time_embedding_dim),
+            act=nn.SiLU(),
+            linear_2=nn.Linear(time_embedding_dim, time_embedding_dim),
+        ))
+
+        # controlnet cond embedding:
+        self.controlnet_cond_embedding = nn.ModuleDict(dict(
+            conv_in=nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            blocks=nn.ModuleList([
+                # 16 -> 32
+                nn.Conv2d(16, 16, kernel_size=3, padding=1),
+                nn.Conv2d(16, 32, kernel_size=3, padding=1, stride=2),
+                # 32 -> 96
+                nn.Conv2d(32, 32, kernel_size=3, padding=1),
+                nn.Conv2d(32, 96, kernel_size=3, padding=1, stride=2),
+                # 96 -> 256
+                nn.Conv2d(96, 96, kernel_size=3, padding=1),
+                nn.Conv2d(96, 256, kernel_size=3, padding=1, stride=2),
+            ]),
+            conv_out=zero_module(nn.Conv2d(256, 320, kernel_size=3, padding=1)),
+        ))
+
+        # actual unet blocks:
+
+        self.conv_in = nn.Conv2d(4, 320, kernel_size=3, padding=1)
+
+        self.down_blocks = nn.ModuleList([
+            # 320 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                    ResnetBlock2D(320, 320, time_embedding_dim),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 320 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(320, 640, time_embedding_dim),
+                    ResnetBlock2D(640, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                downsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)))]),
+            )),
+            # 640 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+            )),
+        ])
+
+        self.controlnet_down_blocks = nn.ModuleList([
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+        ])
+
+        self.mid_block = nn.ModuleDict(dict(
+            resnets=nn.ModuleList([
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+                ResnetBlock2D(1280, 1280, time_embedding_dim),
+            ]),
+            attentions=nn.ModuleList([TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10)]),
+        ))
+
+        self.controlnet_mid_block = zero_module(nn.Conv2d(1280, 1280, kernel_size=1))
+
+        self.up_blocks = nn.ModuleList([
+            # 1280 -> 1280
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(1280 + 1280, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280 + 1280, 1280, time_embedding_dim),
+                    ResnetBlock2D(1280 + 640, 1280, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                    TransformerDecoder2D(1280, encoder_hidden_states_dim, num_transformer_blocks=10),
+                ]),
+                upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(1280, 1280, kernel_size=3, padding=1)))]),
+            )),
+            # 1280 -> 640
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(1280 + 640, 640, time_embedding_dim),
+                    ResnetBlock2D(640 + 640, 640, time_embedding_dim),
+                    ResnetBlock2D(640 + 320, 640, time_embedding_dim),
+                ]),
+                attentions=nn.ModuleList([
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                    TransformerDecoder2D(640, encoder_hidden_states_dim, num_transformer_blocks=2),
+                ]),
+                upsamplers=nn.ModuleList([nn.ModuleDict(dict(conv=nn.Conv2d(640, 640, kernel_size=3, padding=1)))]),
+            )),
+            # 640 -> 320
+            nn.ModuleDict(dict(
+                resnets=nn.ModuleList([
+                    ResnetBlock2D(640 + 320, 320, time_embedding_dim),
+                    ResnetBlock2D(320 + 320, 320, time_embedding_dim),
+                    ResnetBlock2D(320 + 320, 320, time_embedding_dim),
+                ]),
+            ))
+        ])
+
+        # take the output of transformer(resnet(hidden_states)) and project it to
+        # the number of residual channels for the same block
+        self.controlnet_up_blocks = nn.ModuleList([
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 1280, kernel_size=1)),
+            zero_module(nn.Conv2d(1280, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 640, kernel_size=1)),
+            zero_module(nn.Conv2d(640, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+            zero_module(nn.Conv2d(320, 320, kernel_size=1)),
+        ])
+
+        self.conv_norm_out = nn.GroupNorm(32, 320)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(320, 4, kernel_size=3, padding=1)
+
+        self.controlnet_conv_out = zero_module(nn.Conv2d(4, 4, kernel_size=1))
+
+        # fmt: on
+
+    def forward(
+        self,
+        x_t,
+        t,
+        encoder_hidden_states,
+        micro_conditioning,
+        pooled_encoder_hidden_states,
+        controlnet_cond,
+    ):
+        hidden_state = x_t
+
+        t = self.get_sinusoidal_timestep_embedding(t)
+        t = t.to(dtype=hidden_state.dtype)
+        t = self.time_embedding["linear_1"](t)
+        t = self.time_embedding["act"](t)
+        t = self.time_embedding["linear_2"](t)
+
+        additional_conditioning = self.get_sinusoidal_micro_conditioning_embedding(micro_conditioning)
+        additional_conditioning = additional_conditioning.to(dtype=hidden_state.dtype)
+        additional_conditioning = additional_conditioning.flatten(1)
+        additional_conditioning = torch.concat([pooled_encoder_hidden_states, additional_conditioning], dim=-1)
+        additional_conditioning = self.add_embedding["linear_1"](additional_conditioning)
+        additional_conditioning = self.add_embedding["act"](additional_conditioning)
+        additional_conditioning = self.add_embedding["linear_2"](additional_conditioning)
+
+        t = t + additional_conditioning
+
+        controlnet_cond = self.controlnet_cond_embedding["conv_in"](controlnet_cond)
+        controlnet_cond = F.silu(controlnet_cond)
+
+        for block in self.controlnet_cond_embedding["blocks"]:
+            controlnet_cond = F.silu(block(controlnet_cond))
+
+        controlnet_cond = self.controlnet_cond_embedding["conv_out"](controlnet_cond)
+
+        hidden_state = self.conv_in(hidden_state)
+
+        hidden_state = hidden_state + controlnet_cond
+
+        residuals = [hidden_state]
+
+        add_to_down_block_input = self.controlnet_down_blocks[0](hidden_state)
+        add_to_down_block_inputs = [add_to_down_block_input]
+
+        for down_block in self.down_blocks:
+            for i, resnet in enumerate(down_block["resnets"]):
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in down_block:
+                    hidden_state = down_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+                if len(add_to_down_block_inputs) < len(self.controlnet_down_blocks):
+                    add_to_down_block_input = self.controlnet_down_blocks[len(add_to_down_block_inputs)](hidden_state)
+                    add_to_down_block_inputs.append(add_to_down_block_input)
+
+                residuals.append(hidden_state)
+
+            if "downsamplers" in down_block:
+                hidden_state = down_block["downsamplers"][0]["conv"](hidden_state)
+
+                if len(add_to_down_block_inputs) < len(self.controlnet_down_blocks):
+                    add_to_down_block_input = self.controlnet_down_blocks[len(add_to_down_block_inputs)](hidden_state)
+                    add_to_down_block_inputs.append(add_to_down_block_input)
+
+                residuals.append(hidden_state)
+
+        hidden_state = self.mid_block["resnets"][0](hidden_state, t)
+        hidden_state = self.mid_block["attentions"][0](hidden_state, encoder_hidden_states)
+        hidden_state = self.mid_block["resnets"][1](hidden_state, t)
+
+        mid_block_res_sample = self.controlnet_mid_block(hidden_state)
+
+        down_block_res_samples = []
+
+        for up_block in self.up_blocks:
+            for i, resnet in enumerate(up_block["resnets"]):
+                residual = residuals.pop()
+
+                hidden_state = torch.concat([hidden_state, residual], dim=1)
+
+                hidden_state = resnet(hidden_state, t)
+
+                if "attentions" in up_block:
+                    hidden_state = up_block["attentions"][i](hidden_state, encoder_hidden_states)
+
+                down_block_res_sample = self.controlnet_up_blocks[len(down_block_res_samples)](hidden_state)
+                down_block_res_samples.insert(0, down_block_res_sample)
+
+            if "upsamplers" in up_block:
+                hidden_state = F.interpolate(hidden_state, scale_factor=2.0, mode="nearest")
+                hidden_state = up_block["upsamplers"][0]["conv"](hidden_state)
+
+        hidden_state = self.conv_norm_out(hidden_state)
+        hidden_state = self.conv_act(hidden_state)
+        hidden_state = self.conv_out(hidden_state)
+
+        add_to_output = self.controlnet_conv_out(hidden_state)
+
+        return dict(
+            down_block_res_samples=down_block_res_samples,
+            mid_block_res_sample=mid_block_res_sample,
+            add_to_down_block_inputs=add_to_down_block_inputs,
+            add_to_output=add_to_output,
+        )
+
+    @classmethod
+    def from_unet(cls, unet):
+        controlnet = cls()
+
+        controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+        controlnet.add_embedding.load_state_dict(unet.add_embedding.state_dict())
+
+        controlnet.conv_in.load_state_dict(unet.conv_in.state_dict())
+
+        controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())
+        controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())
+        controlnet.up_blocks.load_state_dict(unet.up_blocks.state_dict())
+
+        controlnet.conv_norm_out.load_state_dict(unet.conv_norm_out.state_dict())
+        controlnet.conv_out.load_state_dict(unet.conv_out.state_dict())
+
+        return controlnet
+
+
+class SDXLAdapter(nn.Module, ModelUtils):
+    def __init__(self):
+        super().__init__()
+
+        # fmt: off
+
+        self.adapter = nn.ModuleDict(dict(
+            # 3 -> 768
+            unshuffle=nn.PixelUnshuffle(16),
+
+            # 768 -> 320
+            conv_in=nn.Conv2d(768, 320, kernel_size=3, padding=1),
+
+            body=nn.ModuleList([
+                # 320 -> 320
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList(
+                        nn.ModuleDict(dict(block1=nn.Conv2d(320, 320, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(320, 320, kernel_size=1))),
+                        nn.ModuleDict(dict(block1=nn.Conv2d(320, 320, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(320, 320, kernel_size=1))),
+                    )
+                )),
+                # 320 -> 640
+                nn.ModuleDict(dict(
+                    in_conv=nn.Conv2d(320, 640, kernel_size=1),
+                    resnets=nn.ModuleList(
+                        nn.ModuleDict(dict(block1=nn.Conv2d(640, 640, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(640, 640, kernel_size=1))),
+                        nn.ModuleDict(dict(block1=nn.Conv2d(640, 640, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(640, 640, kernel_size=1))),
+                    )
+                )),
+                # 640 -> 1280
+                nn.ModuleDict(dict(
+                    downsample=nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
+                    in_conv=nn.Conv2d(640, 1280, kernel_size=1),
+                    resnets=nn.ModuleList(
+                        nn.ModuleDict(dict(block1=nn.Conv2d(1280, 1280, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(1280, 1280, kernel_size=1))),
+                        nn.ModuleDict(dict(block1=nn.Conv2d(1280, 1280, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(1280, 1280, kernel_size=1))),
+                    )
+                )),
+                # 1280 -> 1280
+                nn.ModuleDict(dict(
+                    resnets=nn.ModuleList(
+                        nn.ModuleDict(dict(block1=nn.Conv2d(1280, 1280, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(1280, 1280, kernel_size=1))),
+                        nn.ModuleDict(dict(block1=nn.Conv2d(1280, 1280, kernel_size=3, padding=1), act=nn.ReLU(), block2=nn.Conv2d(1280, 1280, kernel_size=1))),
+                    )
+                )),
+            ])
+        ))
+
+        # fmt: on
+
+    def forward(self, x):
+        x = self.unshuffle(x)
+        x = self.conv_in(x)
+
+        features = []
+
+        for block in self.body:
+            if "downsample" in block:
+                x = block["downsample"](x)
+
+            if "in_conv" in block:
+                x = block["in_conv"](x)
+
+            for resnet in block["resnets"]:
+                residual = x
+                x = resnet["block1"](x)
+                x = resnet["act"](x)
+                x = resnet["block2"](x)
+                x = residual + x
+
+            features.append(x)
+
+        return features
+
+
+def get_sinusoidal_embedding(
+    indices: torch.Tensor,
+    embedding_dim: int,
+):
+    half_dim = embedding_dim // 2
+    exponent = -math.log(10000) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=indices.device)
+    exponent = exponent / half_dim
+
+    emb = torch.exp(exponent)
+    emb = indices.unsqueeze(-1).float() * emb
+    emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)
+
+    return emb
+
+
+class ResnetBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels, time_embedding_dim=None, eps=1e-5):
+        super().__init__()
+
+        if time_embedding_dim is not None:
+            self.time_emb_proj = nn.Linear(time_embedding_dim, out_channels)
+        else:
+            self.time_emb_proj = None
+
+        self.norm1 = torch.nn.GroupNorm(32, in_channels, eps=eps)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+
+        self.norm2 = nn.GroupNorm(32, out_channels, eps=eps)
+        self.dropout = nn.Dropout(0.0)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+
+        self.nonlinearity = nn.SiLU()
+
+        if in_channels != out_channels:
+            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        else:
+            self.conv_shortcut = None
+
+    def forward(self, hidden_states, temb=None):
+        residual = hidden_states
+
+        if self.time_emb_proj is not None:
+            assert temb is not None
+            temb = self.nonlinearity(temb)
+            temb = self.time_emb_proj(temb)[:, :, None, None]
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            hidden_states = hidden_states + temb
+
+        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:
+            residual = self.conv_shortcut(residual)
+
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class TransformerDecoder2D(nn.Module):
+    def __init__(self, channels, encoder_hidden_states_dim, num_transformer_blocks):
+        super().__init__()
+
+        self.norm = nn.GroupNorm(32, channels, eps=1e-06)
+        self.proj_in = nn.Linear(channels, channels)
+
+        self.transformer_blocks = nn.ModuleList([TransformerDecoderBlock(channels, encoder_hidden_states_dim) for _ in range(num_transformer_blocks)])
+
+        self.proj_out = nn.Linear(channels, channels)
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        batch_size, channels, height, width = hidden_states.shape
+
+        residual = hidden_states
+
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+        hidden_states = self.proj_in(hidden_states)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states)
+
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2).contiguous()
+
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class TransformerDecoderBlock(nn.Module):
+    def __init__(self, channels, encoder_hidden_states_dim):
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(channels)
+        self.attn1 = Attention(channels, channels)
+
+        self.norm2 = nn.LayerNorm(channels)
+        self.attn2 = Attention(channels, encoder_hidden_states_dim)
+
+        self.norm3 = nn.LayerNorm(channels)
+        self.ff = nn.ModuleDict(dict(net=nn.Sequential(GEGLU(channels, 4 * channels), nn.Dropout(0.0), nn.Linear(4 * channels, channels))))
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states
+
+        hidden_states = self.attn2(self.norm2(hidden_states), encoder_hidden_states) + hidden_states
+
+        hidden_states = self.ff["net"](self.norm3(hidden_states)) + hidden_states
+
+        return hidden_states
+
+
+class Attention(nn.Module):
+    def __init__(self, channels, encoder_hidden_states_dim, qkv_bias=False):
+        super().__init__()
+        self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+        self.to_k = nn.Linear(encoder_hidden_states_dim, channels, bias=qkv_bias)
+        self.to_v = nn.Linear(encoder_hidden_states_dim, channels, bias=qkv_bias)
+        self.to_out = nn.Sequential(nn.Linear(channels, channels), nn.Dropout(0.0))
+
+    def forward(self, hidden_states, encoder_hidden_states=None):
+        batch_size, q_seq_len, channels = hidden_states.shape
+        head_dim = 64
+
+        if encoder_hidden_states is not None:
+            kv = encoder_hidden_states
+        else:
+            kv = hidden_states
+
+        kv_seq_len = kv.shape[1]
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(kv)
+        value = self.to_v(kv)
+
+        query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).contiguous()
+        key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+        value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+
+        hidden_states = xformers.ops.memory_efficient_attention(query, key, value)
+
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = hidden_states.reshape(batch_size, q_seq_len, channels).contiguous()
+
+        hidden_states = self.to_out(hidden_states)
+
+        return hidden_states
+
+
+class GEGLU(nn.Module):
+    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)
+
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module