Upload 28 files

.gitignore

@@ -0,0 +1 @@
+pretrain_weights/

README.md

@@ -1,3 +1,5 @@
+-- LSDM for Crack Segmentation dataset expending
+
 ---
 title: LSDM
 emoji: 💻

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+from evolution import random_walk
+from generate import generate
+
+def process_random_walk(img):
+    img1, _ = random_walk(img)
+
+    return img1
+
+def process_first_generation(img1, model_path="pretrain_weights/b2m/unet_ema"):
+
+    generated_images = generate(img1, model_path)
+    return generated_images[0]
+
+def process_second_generation(img1, model_path="pretrain_weights/m2i/unet_ema"):
+
+    generated_images = generate(img1, model_path)
+    return generated_images[0]
+
+# 创建 Gradio 接口
+with gr.Blocks() as app:
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(value="figs/4.png", image_mode='L', type='numpy', label="Upload Grayscale Image")
+           
+            process_button_1 = gr.Button("1. Process Evolution")
+
+        with gr.Column():
+            output_image_1 = gr.Image(value="figs/4_1.png", image_mode='L', type="numpy", label="After Evolution Image",sources=[])
+            process_button_2 = gr.Button("2. Generate Masks")
+
+    with gr.Row():
+        with gr.Column():
+            output_image_3 = gr.Image(value="figs/4_1_mask.png", image_mode='L', type="numpy", label="Generated Mask Image",sources=[])
+            process_button_3 = gr.Button("3. Generate Images")
+        with gr.Column():
+            output_image_5 = gr.Image(value="figs/4_1.jpg", type="numpy", image_mode='RGB', label="Final Generated Image 1",sources=[])
+    
+
+    process_button_1.click(
+        process_random_walk,
+        inputs=[input_image],
+        outputs=[output_image_1]
+    )
+
+    process_button_2.click(
+        process_first_generation,
+        inputs=[output_image_1],
+        outputs=[output_image_3]
+    )
+
+    process_button_3.click(
+        process_second_generation,
+        inputs=[output_image_3],
+        outputs=[output_image_5]
+    )
+
+app.launch()

diffusion_module/nn.py

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

diffusion_module/unet.py

@@ -0,0 +1,1315 @@
+from abc import abstractmethod
+
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .nn import (
+    SiLU,
+    checkpoint,
+    conv_nd,
+    linear,
+    avg_pool_nd,
+    zero_module,
+    normalization,
+    timestep_embedding,
+    convert_module_to_f16
+)
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.models.modeling_utils import ModelMixin
+from dataclasses import dataclass
+
+@dataclass
+class UNet2DOutput(BaseOutput):
+    """
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Hidden states output. Output of last layer of model.
+    """
+
+    sample: th.FloatTensor
+
+
+class AttentionPool2d(nn.Module):
+    """
+    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+    """
+
+    def __init__(
+        self,
+        spacial_dim: int,
+        embed_dim: int,
+        num_heads_channels: int,
+        output_dim: int = None,
+    ):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(
+            th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5
+        )
+        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+        self.num_heads = embed_dim // num_heads_channels
+        self.attention = QKVAttention(self.num_heads)
+
+    def forward(self, x):
+        b, c, *_spatial = x.shape
+        x = x.reshape(b, c, -1)  # NC(HW)
+        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)
+        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)
+        x = self.qkv_proj(x)
+        x = self.attention(x)
+        x = self.c_proj(x)
+        return x[:, :, 0]
+
+
+class TimestepBlock(nn.Module):
+    """
+    Any module where forward() takes timestep embeddings as a second argument.
+    """
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+
+class CondTimestepBlock(nn.Module):
+    """
+    Any module where forward() takes timestep embeddings as a second argument.
+    """
+
+    @abstractmethod
+    def forward(self, x, cond, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+"""
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock, CondTimestepBlock):
+    
+    def forward(self, x, cond, emb):
+        for layer in self:
+            if isinstance(layer, CondTimestepBlock):
+                x = layer(x, cond, emb)
+            elif isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+        return x
+"""
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock, CondTimestepBlock):
+    def forward(self, x, cond, emb):
+        outputs_list = []  # 创建一个空列表来存储第二个输出
+        for layer in self:
+            if isinstance(layer, CondTimestepBlock):
+                # 调用layer并检查输出是否为一个元组
+                result = layer(x, cond, emb)
+                if isinstance(result, tuple) and len(result) == 2:
+                    x, additional_output = result
+                    outputs_list.append(additional_output)  # 将第二个输出添加到列表
+                else:
+                    x = result
+            elif isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+
+        if outputs_list == []:
+            return x
+        else:
+            return x, outputs_list  # 返回最终的x和所有附加输出的列表
+
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=1
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class SPADEGroupNorm(nn.Module):
+    def __init__(self, norm_nc, label_nc, eps = 1e-5,debug = False):
+        super().__init__()
+        self.debug = debug
+        self.norm = nn.GroupNorm(32, norm_nc, affine=False) # 32/16
+
+        self.eps = eps
+        nhidden = 128
+        self.mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+        self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+    def forward(self, x, segmap):
+        # Part 1. generate parameter-free normalized activations
+        x = self.norm(x)
+        
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+        actv = self.mlp_shared(segmap)
+        gamma = self.mlp_gamma(actv)
+        beta = self.mlp_beta(actv)
+
+        # apply scale and bias
+        if self.debug:
+            return x * (1 + gamma) + beta, (beta.detach().cpu(), gamma.detach().cpu())
+        else:
+            return x * (1 + gamma) + beta
+        
+
+class AdaIN(nn.Module):
+    def __init__(self, num_features):
+        super().__init__()
+        self.instance_norm = th.nn.InstanceNorm2d(num_features, affine=False, track_running_stats=False)
+
+    def forward(self, x, alpha, gamma):
+        assert x.shape[:2] == alpha.shape[:2] == gamma.shape[:2]
+        norm = self.instance_norm(x)
+        return alpha * norm + gamma
+
+class RESAILGroupNorm(nn.Module):
+    def __init__(self, norm_nc, label_nc, guidance_nc, eps = 1e-5):
+        super().__init__()
+
+        self.norm = nn.GroupNorm(32, norm_nc, affine=False) # 32/16
+
+        # SPADE
+        self.eps = eps
+        nhidden = 128
+        self.mask_mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+
+        self.mask_mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.mask_mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+
+        # Guidance
+
+        self.conv_s = th.nn.Conv2d(label_nc, nhidden * 2, 3, 2)
+        self.pool_s = th.nn.AdaptiveAvgPool2d(1)
+        self.conv_s2 = th.nn.Conv2d(nhidden * 2, nhidden * 2, 1, 1)
+
+        self.conv1 = th.nn.Conv2d(guidance_nc, nhidden, 3, 1, padding=1)
+        self.adaIn1 = AdaIN(norm_nc * 2)
+        self.relu1 = nn.ReLU()
+
+        self.conv2 = th.nn.Conv2d(nhidden, nhidden, 3, 1, padding=1)
+        self.adaIn2 = AdaIN(norm_nc * 2)
+        self.relu2 = nn.ReLU()
+        self.conv3 = th.nn.Conv2d(nhidden, nhidden, 3, 1, padding=1)
+
+        self.guidance_mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.guidance_mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+        self.blending_gamma = nn.Parameter(th.zeros(1), requires_grad=True)
+        self.blending_beta = nn.Parameter(th.zeros(1), requires_grad=True)
+        self.norm_nc = norm_nc
+
+    def forward(self, x, segmap, guidance):
+        # Part 1. generate parameter-free normalized activations
+        x = self.norm(x)
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+        mask_actv = self.mask_mlp_shared(segmap)
+        mask_gamma = self.mask_mlp_gamma(mask_actv)
+        mask_beta = self.mask_mlp_beta(mask_actv)
+
+
+        # Part 3. produce scaling and bias conditioned on feature guidance
+        guidance = F.interpolate(guidance, size=x.size()[2:], mode='bilinear')
+
+        f_s_1 = self.conv_s(segmap)
+        c1 = self.pool_s(f_s_1)
+        c2 = self.conv_s2(c1)
+
+        f1 = self.conv1(guidance)
+
+        f1 = self.adaIn1(f1, c1[:, : 128, ...], c1[:, 128:, ...])
+        f2 = self.relu1(f1)
+
+        f2 = self.conv2(f2)
+        f2 = self.adaIn2(f2, c2[:, : 128, ...], c2[:, 128:, ...])
+        f2 = self.relu2(f2)
+        guidance_actv = self.conv3(f2)
+
+        guidance_gamma = self.guidance_mlp_gamma(guidance_actv)
+        guidance_beta = self.guidance_mlp_beta(guidance_actv)
+
+        gamma_alpha = F.sigmoid(self.blending_gamma)
+        beta_alpha = F.sigmoid(self.blending_beta)
+
+        gamma_final = gamma_alpha * guidance_gamma + (1 - gamma_alpha) * mask_gamma
+        beta_final = beta_alpha * guidance_beta + (1 - beta_alpha) * mask_beta
+        out = x * (1 + gamma_final) + beta_final
+
+        # apply scale and bias
+        return out
+
+class SPMGroupNorm(nn.Module):
+    def __init__(self, norm_nc, label_nc, feature_nc, eps = 1e-5):
+        super().__init__()
+        print("use SPM")
+
+        self.norm = nn.GroupNorm(32, norm_nc, affine=False) # 32/16
+
+        # SPADE
+        self.eps = eps
+        nhidden = 128
+        self.mask_mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+
+        self.mask_mlp_gamma1 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.mask_mlp_beta1 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+        self.mask_mlp_gamma2 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.mask_mlp_beta2 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+
+        # Feature
+        self.feature_mlp_shared = nn.Sequential(
+            nn.Conv2d(feature_nc, nhidden, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+
+        self.feature_mlp_gamma1 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.feature_mlp_beta1 = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+
+    def forward(self, x, segmap, guidance):
+        # Part 1. generate parameter-free normalized activations
+        x = self.norm(x)
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+        mask_actv = self.mask_mlp_shared(segmap)
+        mask_gamma1 = self.mask_mlp_gamma1(mask_actv)
+        mask_beta1 = self.mask_mlp_beta1(mask_actv)
+
+        mask_gamma2 = self.mask_mlp_gamma2(mask_actv)
+        mask_beta2 = self.mask_mlp_beta2(mask_actv)
+
+
+        # Part 3. produce scaling and bias conditioned on feature guidance
+        guidance = F.interpolate(guidance, size=x.size()[2:], mode='bilinear')
+        feature_actv = self.feature_mlp_shared(guidance)
+        feature_gamma1 = self.feature_mlp_gamma1(feature_actv)
+        feature_beta1 = self.feature_mlp_beta1(feature_actv)
+
+        gamma_final = feature_gamma1 * (1 + mask_gamma1) + mask_beta1
+        beta_final = feature_beta1 * (1 + mask_gamma2) + mask_beta2
+
+        out = x * (1 + gamma_final) + beta_final
+
+        # apply scale and bias
+        return out
+
+
+class ResBlock(TimestepBlock):
+    """
+    A residual block that can optionally change the number of channels.
+
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+
+        return th.utils.checkpoint.checkpoint(self._forward, x ,emb)
+        # return checkpoint(
+        #     self._forward, (x, emb), self.parameters(), self.use_checkpoint
+        # )
+
+    def _forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb)#.type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = th.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+class SDMResBlock(CondTimestepBlock):
+    """
+    A residual block that can optionally change the number of channels.
+
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        c_channels=3,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+        SPADE_type = "spade",
+        guidance_nc = None,
+        debug = False
+    ):
+        super().__init__()
+        self.channels = channels
+        self.guidance_nc = guidance_nc
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.SPADE_type = SPADE_type
+        self.debug = debug
+        if self.SPADE_type == "spade":
+            self.in_norm = SPADEGroupNorm(channels, c_channels, debug=self.debug)
+        elif self.SPADE_type == "RESAIL":
+            self.in_norm = RESAILGroupNorm(channels, c_channels, guidance_nc)
+        elif self.SPADE_type == "SPM":
+            self.in_norm = SPMGroupNorm(channels, c_channels, guidance_nc)
+        self.in_layers = nn.Sequential(
+            SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+        
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+
+        if self.SPADE_type == "spade":
+            self.out_norm = SPADEGroupNorm(self.out_channels, c_channels,debug=self.debug)
+        elif self.SPADE_type == "RESAIL":
+            self.out_norm = RESAILGroupNorm(self.out_channels, c_channels, guidance_nc)
+        elif self.SPADE_type == "SPM":
+            self.out_norm = SPMGroupNorm(self.out_channels, c_channels, guidance_nc)
+
+        self.out_layers = nn.Sequential(
+            SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, cond, emb):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        return th.utils.checkpoint.checkpoint(self._forward, x, cond, emb)
+        # return checkpoint(
+        #     self._forward, (x, cond, emb), self.parameters(), self.use_checkpoint
+        # )
+
+    def _forward(self, x, cond, emb):
+        if self.SPADE_type == "RESAIL" or self.SPADE_type == "SPM":
+            assert self.guidance_nc is not None, "Please set guidance_nc when you use RESAIL"
+            guidance = x[: ,x.shape[1] - self.guidance_nc:, ...]
+        else:
+            guidance = None
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            if self.SPADE_type == "spade":
+                if not self.debug:
+                    h = self.in_norm(x, cond)
+                else:
+                    h, (b1,g1) = self.in_norm(x, cond)
+            elif self.SPADE_type == "RESAIL" or self.SPADE_type == "SPM":
+                h = self.in_norm(x, cond, guidance)
+
+            h = in_rest(h)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            if self.SPADE_type == "spade":
+                if not self.debug:
+                    h = self.in_norm(x, cond)
+                else:
+                    h, (b1,g1) = self.in_norm(x, cond)
+            elif self.SPADE_type == "RESAIL" or self.SPADE_type == "SPM":
+                h = self.in_norm(x, cond, guidance)
+            h = self.in_layers(h)
+
+        emb_out = self.emb_layers(emb)#.type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            scale, shift = th.chunk(emb_out, 2, dim=1)
+            if self.SPADE_type == "spade":
+                if not self.debug:
+                    h = self.out_norm(h, cond)
+                else:
+                    h, (b2,g2) = self.out_norm(h, cond)
+            elif self.SPADE_type == "RESAIL" or self.SPADE_type == "SPM":
+                h = self.out_norm(h, cond, guidance)
+
+            h = h * (1 + scale) + shift
+            h = self.out_layers(h)
+        else:
+            h = h + emb_out
+            if self.SPADE_type == "spade":
+                h = self.out_norm(h, cond)
+            elif self.SPADE_type == "RESAIL" or self.SPADE_type == "SPM":
+                h = self.out_norm(x, cond, guidance)
+
+            h = self.out_layers(h)
+        if self.debug:
+            extra = {(b1,g1),(b2,g2)}
+            return self.skip_connection(x) + h, extra
+        else:
+            return self.skip_connection(x) + h
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        use_checkpoint=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.use_checkpoint = use_checkpoint
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, channels * 3, 1)
+        if use_new_attention_order:
+            # split qkv before split heads
+            self.attention = QKVAttention(self.num_heads)
+        else:
+            # split heads before split qkv
+            self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+    def forward(self, x):
+        return th.utils.checkpoint.checkpoint(self._forward, x)
+        #return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint)
+
+    def _forward(self, x):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+def count_flops_attn(model, _x, y):
+    """
+    A counter for the `thop` package to count the operations in an
+    attention operation.
+    Meant to be used like:
+        macs, params = thop.profile(
+            model,
+            inputs=(inputs, timestamps),
+            custom_ops={QKVAttention: QKVAttention.count_flops},
+        )
+    """
+    b, c, *spatial = y[0].shape
+    num_spatial = int(np.prod(spatial))
+    # We perform two matmuls with the same number of ops.
+    # The first computes the weight matrix, the second computes
+    # the combination of the value vectors.
+    matmul_ops = 2 * b * (num_spatial ** 2) * c
+    model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+    """
+    A module which performs QKV attention and splits in a different order.
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.chunk(3, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class UNetModel(ModelMixin, ConfigMixin):
+    """
+    The full UNet model with attention and timestep embedding.
+
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    """
+
+    _supports_gradient_checkpointing = True
+    @register_to_config
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        use_checkpoint=False,
+        use_fp16=True,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        mask_emb="resize",
+        SPADE_type="spade",
+        debug = False
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.sample_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        self.debug = debug
+
+        self.mask_emb = mask_emb
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        ch = input_ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] #ch=256
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                #print(ds)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+        self.middle_block = TimestepEmbedSequential(
+            SDMResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                c_channels=num_classes if mask_emb == "resize" else num_classes*4,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            SDMResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                c_channels=num_classes if mask_emb == "resize" else num_classes*4 ,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                #print(ch, ich)
+                layers = [
+                    SDMResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        c_channels=num_classes if mask_emb == "resize" else num_classes*4,
+                        out_channels=int(model_channels * mult),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        SPADE_type=SPADE_type,
+                        guidance_nc = ich,
+                        debug=self.debug,
+                    )
+                ]
+                ch = int(model_channels * mult)
+                #print(ds)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads_upsample,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    layers.append(
+                        SDMResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            c_channels=num_classes if mask_emb == "resize" else num_classes*4,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                            debug=self.debug
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            SiLU(),
+            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
+        )
+    def _set_gradient_checkpointing(self, module, value=False):
+        #if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):
+        module.gradient_checkpointing = value
+    def forward(self, x, y=None, timesteps=None ):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+
+        hs = []
+        if not th.is_tensor(timesteps):
+            timesteps = th.tensor([timesteps], dtype=th.long, device=x.device)
+        elif th.is_tensor(timesteps) and len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(x.device)
+
+        timesteps = timestep_embedding(timesteps, self.model_channels).type(x.dtype).to(x.device)
+        emb = self.time_embed(timesteps)
+
+        y = y.type(self.dtype)
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            # input_blocks have no any opts for y
+            h = module(h, y, emb)
+            #print(h.shape)
+            hs.append(h)
+
+        h = self.middle_block(h, y, emb)
+
+        if self.debug:
+            extra_list = []
+
+        for module in self.output_blocks:
+            temp = hs.pop()
+
+            #print("before:", h.shape, temp.shape)
+            # copy padding to match the downsample size
+            if h.shape[2] != temp.shape[2]:
+                p1d = (0, 0, 0, 1)
+                h = F.pad(h, p1d, "replicate")
+
+            if h.shape[3] != temp.shape[3]:
+                p2d = (0, 1, 0, 0)
+                h = F.pad(h, p2d, "replicate")
+            #print("after:", h.shape, temp.shape)
+
+            h = th.cat([h, temp], dim=1)
+            if self.debug:
+                h, extra = module(h, y, emb)
+                extra_list.append(extra)
+            else:
+                h = module(h, y, emb)
+
+        h = h.type(x.dtype)
+
+        if not self.debug:
+            return UNet2DOutput(sample=self.out(h))
+        else:
+            return UNet2DOutput(sample=self.out(h)), extra_list
+
+
+class SuperResModel(UNetModel):
+    """
+    A UNetModel that performs super-resolution.
+
+    Expects an extra kwarg `low_res` to condition on a low-resolution image.
+    """
+
+    def __init__(self, image_size, in_channels, *args, **kwargs):
+        super().__init__(image_size, in_channels * 2, *args, **kwargs)
+
+    def forward(self, x, cond, timesteps, low_res=None, **kwargs):
+        _, _, new_height, new_width = x.shape
+        upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
+        x = th.cat([x, upsampled], dim=1)
+        return super().forward(x, cond, timesteps, **kwargs)
+
+
+class EncoderUNetModel(nn.Module):
+    """
+    The half UNet model with attention and timestep embedding.
+
+    For usage, see UNet.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        pool="adaptive",
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        ch = int(channel_mult[0] * model_channels)
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
+        )
+        self._feature_size = ch
+        input_block_chans = [ch]
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=int(mult * model_channels),
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = int(mult * model_channels)
+                if ds in attention_resolutions:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+        self.pool = pool
+        if pool == "adaptive":
+            self.out = nn.Sequential(
+                normalization(ch),
+                SiLU(),
+                nn.AdaptiveAvgPool2d((1, 1)),
+                zero_module(conv_nd(dims, ch, out_channels, 1)),
+                nn.Flatten(),
+            )
+        elif pool == "attention":
+            assert num_head_channels != -1
+            self.out = nn.Sequential(
+                normalization(ch),
+                SiLU(),
+                AttentionPool2d(
+                    (image_size // ds), ch, num_head_channels, out_channels
+                ),
+            )
+        elif pool == "spatial":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                nn.ReLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        elif pool == "spatial_v2":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                normalization(2048),
+                SiLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        else:
+            raise NotImplementedError(f"Unexpected {pool} pooling")
+    def forward(self, x, timesteps):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :return: an [N x K] Tensor of outputs.
+        """
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        results = []
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            if self.pool.startswith("spatial"):
+                results.append(h.type(x.dtype).mean(dim=(2, 3)))
+        h = self.middle_block(h, emb)
+        if self.pool.startswith("spatial"):
+            results.append(h.type(x.dtype).mean(dim=(2, 3)))
+            h = th.cat(results, axis=-1)
+            return self.out(h)
+        else:
+            h = h.type(x.dtype)
+            return self.out(h)
+

diffusion_module/unet_2d_sdm.py

@@ -0,0 +1,357 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.models.embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from .unet_2d_blocks import UNetSDMMidBlock2D, get_down_block, get_up_block, UNetSDMMidBlock2D
+from diffusers.loaders import UNet2DConditionLoadersMixin
+
+@dataclass
+class UNet2DOutput(BaseOutput):
+    """
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Hidden states output. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor
+
+class SDMUNet2DModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
+    r"""
+    UNet2DModel is a 2D UNet model that takes in a noisy sample and a timestep and returns sample shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library
+    implements for all the model (such as downloading or saving, etc.)
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample. Dimensions must be a multiple of `2 ** (len(block_out_channels) -
+            1)`.
+        in_channels (`int`, *optional*, defaults to 3): Number of channels in the input image.
+        out_channels (`int`, *optional*, defaults to 3): Number of channels in the output.
+        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
+        time_embedding_type (`str`, *optional*, defaults to `"positional"`): Type of time embedding to use.
+        freq_shift (`int`, *optional*, defaults to 0): Frequency shift for fourier time embedding.
+        flip_sin_to_cos (`bool`, *optional*, defaults to :
+            obj:`True`): Whether to flip sin to cos for fourier time embedding.
+        down_block_types (`Tuple[str]`, *optional*, defaults to :
+            obj:`("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`): Tuple of downsample block
+            types.
+        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`):
+            The mid block type. Choose from `UNetMidBlock2D` or `UnCLIPUNetMidBlock2D`.
+        up_block_types (`Tuple[str]`, *optional*, defaults to :
+            obj:`("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`): Tuple of upsample block types.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to :
+            obj:`(224, 448, 672, 896)`): Tuple of block output channels.
+        layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block.
+        mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block.
+        downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.
+        norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for the normalization.
+        norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for the normalization.
+        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
+            for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`.
+        class_embed_type (`str`, *optional*, defaults to None):
+            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
+            `"timestep"`, or `"identity"`.
+        num_class_embeds (`int`, *optional*, defaults to None):
+            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
+            class conditioning with `class_embed_type` equal to `None`.
+    """
+
+    _supports_gradient_checkpointing = True
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[Union[int, Tuple[int, int]]] = None,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        center_input_sample: bool = False,
+        time_embedding_type: str = "positional",
+        freq_shift: int = 0,
+        flip_sin_to_cos: bool = True,
+        down_block_types: Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
+        up_block_types: Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
+        block_out_channels: Tuple[int] = (224, 448, 672, 896),
+        layers_per_block: int = 2,
+        mid_block_scale_factor: float = 1,
+        downsample_padding: int = 1,
+        act_fn: str = "silu",
+        attention_head_dim: Optional[int] = 8,
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-5,
+        resnet_time_scale_shift: str = "scale_shift",
+        add_attention: bool = True,
+        class_embed_type: Optional[str] = None,
+        num_class_embeds: Optional[int] = None,
+        segmap_channels: int = 34,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+        self.segmap_channels = segmap_channels
+        time_embed_dim = block_out_channels[0] * 4
+
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))
+
+        # time
+        if time_embedding_type == "fourier":
+            self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16)
+            timestep_input_dim = 2 * block_out_channels[0]
+        elif time_embedding_type == "positional":
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        # class embedding
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attn_num_head_channels=attention_head_dim,
+                downsample_padding=downsample_padding,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetSDMMidBlock2D(
+            in_channels=block_out_channels[-1],
+            temb_channels=time_embed_dim,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            attn_num_head_channels=attention_head_dim,
+            resnet_groups=norm_num_groups,
+            add_attention=add_attention,
+            segmap_channels=segmap_channels,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            is_final_block = i == len(block_out_channels) - 1
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=layers_per_block + 1,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=time_embed_dim,
+                add_upsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attn_num_head_channels=attention_head_dim,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                segmap_channels=segmap_channels,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32)
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=num_groups_out, eps=norm_eps)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1)
+    def _set_gradient_checkpointing(self, module, value=False):
+        #if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):
+        module.gradient_checkpointing = value
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        segmap: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        class_labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DOutput, Tuple]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
+            timestep (`torch.FloatTensor` or `float` or `int): (batch) timesteps
+            class_labels (`torch.FloatTensor`, *optional*, defaults to `None`):
+                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_2d.UNet2DOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.unet_2d.UNet2DOutput`] or `tuple`: [`~models.unet_2d.UNet2DOutput`] if `return_dict` is True,
+            otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.
+        """
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        #print(timestep.shape)
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)
+        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps * torch.ones(sample.shape[0], dtype=timesteps.dtype, device=timesteps.device)
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+        emb = self.time_embedding(t_emb)
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when doing class conditioning")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
+
+        # 2. pre-process
+        skip_sample = sample
+        sample = self.conv_in(sample)
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "skip_conv"):
+                sample, res_samples, skip_sample = downsample_block(
+                    hidden_states=sample, temb=emb, skip_sample=skip_sample,segmap=segmap
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        sample = self.mid_block(sample, segmap, emb)
+
+        # 5. up
+        skip_sample = None
+        for upsample_block in self.up_blocks:
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            if hasattr(upsample_block, "skip_conv"):
+                sample, skip_sample = upsample_block(sample, segmap, res_samples, emb, skip_sample)
+            else:
+                sample = upsample_block(sample, segmap, res_samples, emb)
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if skip_sample is not None:
+            sample += skip_sample
+
+        if self.config.time_embedding_type == "fourier":
+            timesteps = timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]))))
+            sample = sample / timesteps
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DOutput(sample=sample)
+
+
+if __name__ == "__main__":
+    path = 'output.txt'
+    f = open(path, 'w')
+
+    unet = SDMUNet2DModel(
+        sample_size=270,
+        in_channels=3,
+        out_channels=3,
+        layers_per_block=2,
+        block_out_channels=(256, 256, 512, 1024, 1024),
+        down_block_types=(
+            "ResnetDownsampleBlock2D",
+            "ResnetDownsampleBlock2D",
+            "ResnetDownsampleBlock2D",
+            "AttnDownBlock2D",
+            "AttnDownBlock2D",
+        ),
+        up_block_types=(
+            "SDMAttnUpBlock2D",
+            "SDMAttnUpBlock2D",
+            "SDMResnetUpsampleBlock2D",
+            "SDMResnetUpsampleBlock2D",
+            "SDMResnetUpsampleBlock2D",
+        ),
+        segmap_channels=34+1
+    )
+
+    print(unet,file=f)
+    f.close()
+
+    #summary(unet, [(1, 3, 270, 360), (1, 3, 270, 360), (2,)], device="cpu")

diffusion_module/utils/LSDMPipeline_expandDataset.py

@@ -0,0 +1,179 @@
+import inspect
+from typing import List, Optional, Tuple, Union
+
+import torch
+
+from diffusers.models import UNet2DModel, VQModel
+from diffusers.schedulers import DDIMScheduler
+from diffusers.utils import randn_tensor
+from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+import copy
+
+class SDMLDMPipeline(DiffusionPipeline):
+    r"""
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Parameters:
+        vae ([`VQModel`]):
+            Vector-quantized (VQ) Model to encode and decode images to and from latent representations.
+        unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            [`DDIMScheduler`] is to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    def __init__(self, vae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler, torch_dtype=torch.float16, resolution=512, resolution_type="city"):
+        super().__init__()
+        self.register_modules(vae=vae, unet=unet, scheduler=scheduler)
+        self.torch_dtype = torch_dtype
+        self.resolution = resolution
+        self.resolution_type = resolution_type
+    @torch.no_grad()
+    def __call__(
+        self,
+        segmap = None,
+        batch_size: int = 8,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        eta: float = 0.0,
+        num_inference_steps: int = 1000,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        every_step_save: int = None,
+        s: int = 1,
+        num_evolution_per_mask = 10,
+        debug = False,
+        **kwargs,
+    ) -> Union[Tuple, ImagePipelineOutput]:
+        r"""
+        Args:
+            batch_size (`int`, *optional*, defaults to 1):
+                Number of images to generate.
+            generator (`torch.Generator`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`: [`~pipelines.model.ImagePipelineOutput`] if `return_dict` is
+            True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images.
+        """
+        # self.unet.config.sample_size = (64, 64) # (135,180)
+        # self.unet.config.sample_size = (135,180)
+        if self.resolution_type == "crack":
+            self.unet.config.sample_size = (64,64)
+        elif self.resolution_type == "crack_256":
+            self.unet.config.sample_size = (256,256)
+        else:
+            sc = 1080 // self.resolution
+            latent_size = (self.resolution // 4, 1440 // (sc*4))
+            self.unet.config.sample_size = latent_size
+        # 
+        if not isinstance(self.unet.config.sample_size, tuple):
+            self.unet.config.sample_size = (self.unet.config.sample_size, self.unet.config.sample_size)
+
+        if segmap is None:
+            print("Didn't inpute any segmap, use the empty as the input")
+            segmap = torch.zeros(batch_size,self.unet.config.segmap_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1])
+        segmap = segmap.to(self.device).type(self.torch_dtype)
+        if batch_size == 1 and num_evolution_per_mask > batch_size:
+            latents = randn_tensor(
+                (num_evolution_per_mask, self.unet.config.in_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1]),
+                generator=generator,
+        )   
+        else: 
+            latents = randn_tensor(
+                (batch_size, self.unet.config.in_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1]),
+                generator=generator,
+            )
+        latents = latents.to(self.device).type(self.torch_dtype)
+
+        # scale the initial noise by the standard deviation required by the scheduler (need to check)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps)
+
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+
+        extra_kwargs = {}
+        if accepts_eta:
+            extra_kwargs["eta"] = eta
+
+        step_latent = []
+        learn_sigma = True if hasattr(self.scheduler, "variance_type") else False
+        if debug:
+            extra_list_list = []
+            self.unet.debug=True
+        for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
+    
+            latent_model_input = self.scheduler.scale_model_input(latents, t)
+            # predict the noise residual
+            if debug:
+                output, extra_list = self.unet(latent_model_input, segmap, t)   
+                noise_prediction = output.sample
+                extra_list_list.append(extra_list)
+            else:
+                noise_prediction = self.unet(latent_model_input, segmap, t).sample
+            # compute the previous noisy sample x_t -> x_t-1
+            
+
+            if learn_sigma and "learn" in self.scheduler.variance_type:
+                model_pred, var_pred = torch.split(noise_prediction, latents.shape[1], dim=1)
+            else:
+                model_pred = noise_prediction
+            if s > 1.0:
+                if debug:
+                    model_output_zero = self.unet(latent_model_input, torch.zeros_like(segmap), t)[0].sample
+                else:
+                    model_output_zero = self.unet(latent_model_input, torch.zeros_like(segmap), t).sample
+                if learn_sigma and "learn" in self.scheduler.variance_type:
+                    model_output_zero,_ = torch.split(model_output_zero, latents.shape[1], dim=1)
+                model_pred = model_pred + s * (model_pred - model_output_zero)
+                if learn_sigma and "learn" in self.scheduler.variance_type:
+                    recombined = torch.cat((model_pred, var_pred), dim=1)
+            # when apply different scheduler, mean only !!
+            if learn_sigma and "learn" in self.scheduler.variance_type:
+                latents = self.scheduler.step(recombined, t, latents, **extra_kwargs).prev_sample
+            else:
+                latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample
+
+            if every_step_save is not None:
+                if (i+1) % every_step_save == 0:
+                    step_latent.append(copy.deepcopy(latents))
+            
+        if debug:
+            return extra_list_list[-1]
+
+        # decode the image latents with the VAE
+        if every_step_save is not None:
+            image = []
+            for i, l in enumerate(step_latent):
+                l /= self.vae.config.scaling_factor  # (0.18215)
+                #latents /= 7.706491063029163
+                l = self.vae.decode(l, segmap)
+                l = (l / 2 + 0.5).clamp(0, 1)
+                l = l.cpu().permute(0, 2, 3, 1).numpy()
+                if output_type == "pil":
+                    l = self.numpy_to_pil(l)
+                image.append(l)
+        else:
+            latents /= self.vae.config.scaling_factor#(0.18215)
+            #latents /= 7.706491063029163
+            # image = self.vae.decode(latents, segmap).sample
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)

diffusion_module/utils/Pipline.py

@@ -0,0 +1,361 @@
+import inspect
+from typing import List, Optional, Tuple, Union
+
+import torch
+
+from diffusers.models import UNet2DModel, VQModel
+from diffusers.schedulers import DDIMScheduler
+from diffusers.utils import randn_tensor
+from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+import copy
+
+class LDMPipeline(DiffusionPipeline):
+    r"""
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Parameters:
+        vae ([`VQModel`]):
+            Vector-quantized (VQ) Model to encode and decode images to and from latent representations.
+        unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            [`DDIMScheduler`] is to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    def __init__(self, vae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler, torch_dtype=torch.float16):
+        super().__init__()
+        self.register_modules(vae=vae, unet=unet, scheduler=scheduler)
+        self.torch_dtype = torch_dtype
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        batch_size: int = 8,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        eta: float = 0.0,
+        num_inference_steps: int = 1000,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[Tuple, ImagePipelineOutput]:
+        r"""
+        Args:
+            batch_size (`int`, *optional*, defaults to 1):
+                Number of images to generate.
+            generator (`torch.Generator`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`: [`~pipelines.model.ImagePipelineOutput`] if `return_dict` is
+            True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images.
+        """
+        if not isinstance(self.unet.config.sample_size,tuple):
+            self.unet.config.sample_size = (self.unet.config.sample_size,self.unet.config.sample_size)
+
+        latents = randn_tensor(
+            (batch_size, self.unet.config.in_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1]),
+            generator=generator,
+        )
+        latents = latents.to(self.device).type(self.torch_dtype)
+
+        # scale the initial noise by the standard deviation required by the scheduler (need to check)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+
+        extra_kwargs = {}
+        if accepts_eta:
+            extra_kwargs["eta"] = eta
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            latent_model_input = self.scheduler.scale_model_input(latents, t)
+            # predict the noise residual
+            noise_prediction = self.unet(latent_model_input, t).sample
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample
+
+        # decode the image latents with the VAE
+        latents /= self.vae.config.scaling_factor#(0.18215)
+        image = self.vae.decode(latents).sample
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).numpy()
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+
+class SDMLDMPipeline(DiffusionPipeline):
+    r"""
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Parameters:
+        vae ([`VQModel`]):
+            Vector-quantized (VQ) Model to encode and decode images to and from latent representations.
+        unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            [`DDIMScheduler`] is to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    def __init__(self, vae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler, torch_dtype=torch.float16, resolution=512, resolution_type="city"):
+        super().__init__()
+        self.register_modules(vae=vae, unet=unet, scheduler=scheduler)
+        self.torch_dtype = torch_dtype
+        self.resolution = resolution
+        self.resolution_type = resolution_type
+    @torch.no_grad()
+    def __call__(
+        self,
+        segmap = None,
+        batch_size: int = 8,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        eta: float = 0.0,
+        num_inference_steps: int = 1000,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        every_step_save: int = None,
+        s: int = 1,
+        **kwargs,
+    ) -> Union[Tuple, ImagePipelineOutput]:
+        r"""
+        Args:
+            batch_size (`int`, *optional*, defaults to 1):
+                Number of images to generate.
+            generator (`torch.Generator`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`: [`~pipelines.model.ImagePipelineOutput`] if `return_dict` is
+            True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images.
+        """
+        # self.unet.config.sample_size = (64, 64) # (135,180)
+        # self.unet.config.sample_size = (135,180)
+        if self.resolution_type == "crack":
+            self.unet.config.sample_size = (64,64)
+        elif self.resolution_type == "crack_256":
+            self.unet.config.sample_size = (256,256)
+        else:
+            sc = 1080 // self.resolution
+            latent_size = (self.resolution // 4, 1440 // (sc*4))
+            self.unet.config.sample_size = latent_size
+        # 
+        if not isinstance(self.unet.config.sample_size, tuple):
+            self.unet.config.sample_size = (self.unet.config.sample_size, self.unet.config.sample_size)
+
+        if segmap is None:
+            print("Didn't inpute any segmap, use the empty as the input")
+            segmap = torch.zeros(batch_size,self.unet.config.segmap_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1])
+        segmap = segmap.to(self.device).type(self.torch_dtype)
+        latents = randn_tensor(
+            (batch_size, self.unet.config.in_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1]),
+            generator=generator,
+        )
+        latents = latents.to(self.device).type(self.torch_dtype)
+
+        # scale the initial noise by the standard deviation required by the scheduler (need to check)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps)
+
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+
+        extra_kwargs = {}
+        if accepts_eta:
+            extra_kwargs["eta"] = eta
+
+        step_latent = []
+        learn_sigma = True if hasattr(self.scheduler, "variance_type") else False
+        for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
+    
+            latent_model_input = self.scheduler.scale_model_input(latents, t)
+            # predict the noise residual
+            noise_prediction = self.unet(latent_model_input, segmap, t).sample
+            # compute the previous noisy sample x_t -> x_t-1
+            
+
+            if learn_sigma and "learn" in self.scheduler.variance_type:
+                model_pred, var_pred = torch.split(noise_prediction, latents.shape[1], dim=1)
+            else:
+                model_pred = noise_prediction
+            if s > 1.0:
+                model_output_zero = self.unet(latent_model_input, torch.zeros_like(segmap), t).sample
+                if learn_sigma and "learn" in self.scheduler.variance_type:
+                    model_output_zero,_ = torch.split(model_output_zero, latents.shape[1], dim=1)
+                model_pred = model_pred + s * (model_pred - model_output_zero)
+                if learn_sigma and "learn" in self.scheduler.variance_type:
+                    recombined = torch.cat((model_pred, var_pred), dim=1)
+            # when apply different scheduler, mean only !!
+            if learn_sigma and "learn" in self.scheduler.variance_type:
+                latents = self.scheduler.step(recombined, t, latents, **extra_kwargs).prev_sample
+            else:
+                latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample
+
+            if every_step_save is not None:
+                if (i+1) % every_step_save == 0:
+                    step_latent.append(copy.deepcopy(latents))
+
+        # decode the image latents with the VAE
+        if every_step_save is not None:
+            image = []
+            for i, l in enumerate(step_latent):
+                l /= self.vae.config.scaling_factor  # (0.18215)
+                #latents /= 7.706491063029163
+                l = self.vae.decode(l, segmap)
+                l = (l / 2 + 0.5).clamp(0, 1)
+                l = l.cpu().permute(0, 2, 3, 1).numpy()
+                if output_type == "pil":
+                    l = self.numpy_to_pil(l)
+                image.append(l)
+        else:
+            latents /= self.vae.config.scaling_factor#(0.18215)
+            #latents /= 7.706491063029163
+            # image = self.vae.decode(latents, segmap).sample
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+
+class SDMPipeline(DiffusionPipeline):
+    r"""
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Parameters:
+        vae ([`VQModel`]):
+            Vector-quantized (VQ) Model to encode and decode images to and from latent representations.
+        unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            [`DDIMScheduler`] is to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    def __init__(self, unet: UNet2DModel, scheduler: DDIMScheduler, torch_dtype=torch.float16, vae=None):
+        super().__init__()
+        self.register_modules(unet=unet, scheduler=scheduler)
+        self.torch_dtype = torch_dtype
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        segmap = None,
+        batch_size: int = 8,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        eta: float = 0.0,
+        num_inference_steps: int = 1000,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        s: int = 1,
+        **kwargs,
+    ) -> Union[Tuple, ImagePipelineOutput]:
+        r"""
+        Args:
+            batch_size (`int`, *optional*, defaults to 1):
+                Number of images to generate.
+            generator (`torch.Generator`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`: [`~pipelines.model.ImagePipelineOutput`] if `return_dict` is
+            True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images.
+        """
+        self.unet.config.sample_size = (270,360)
+        if not isinstance(self.unet.config.sample_size, tuple):
+            self.unet.config.sample_size = (self.unet.config.sample_size, self.unet.config.sample_size)
+
+        if segmap is None:
+            print("Didn't inpute any segmap, use the empty as the input")
+            segmap = torch.zeros(batch_size,self.unet.config.segmap_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1])
+        segmap = segmap.to(self.device).type(self.torch_dtype)
+        latents = randn_tensor(
+            (batch_size, self.unet.config.in_channels, self.unet.config.sample_size[0], self.unet.config.sample_size[1]),
+            generator=generator,
+        )
+
+        latents = latents.to(self.device).type(self.torch_dtype)
+
+        # scale the initial noise by the standard deviation required by the scheduler (need to check)
+        latents = latents * self.scheduler.init_noise_sigma
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+
+        extra_kwargs = {}
+        if accepts_eta:
+            extra_kwargs["eta"] = eta
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            latent_model_input = self.scheduler.scale_model_input(latents, t)
+            # predict the noise residual
+            noise_prediction = self.unet(latent_model_input, segmap, t).sample
+
+            #noise_prediction = noise_prediction[]
+
+            if s > 1.0:
+                model_output_zero = self.unet(latent_model_input, torch.zeros_like(segmap), t).sample
+                noise_prediction[:, :3] = model_output_zero[:, :3] + s * (noise_prediction[:, :3] - model_output_zero[:, :3])
+
+            #noise_prediction = noise_prediction[:, :3]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            #breakpoint()
+            latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample
+
+        # decode the image latents with the VAE
+        # latents /= self.vae.config.scaling_factor#(0.18215)
+        # image = self.vae.decode(latents).sample
+        image = latents
+        #image = (image + 1) / 2.0
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).numpy()
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+

diffusion_module/utils/loss.py

@@ -0,0 +1,149 @@
+"""
+Helpers for various likelihood-based losses. These are ported from the original
+Ho et al. diffusion models codebase:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py
+"""
+
+import numpy as np
+
+import torch as th
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+
+def variance_KL_loss(latents, noisy_latents, timesteps, model_pred_mean, model_pred_var, noise_scheduler,posterior_mean_coef1, posterior_mean_coef2, posterior_log_variance_clipped):
+    model_pred_mean = model_pred_mean.detach()
+    true_mean = (
+            posterior_mean_coef1.to(device=timesteps.device)[timesteps].float()[..., None, None, None] * latents
+            + posterior_mean_coef2.to(device=timesteps.device)[timesteps].float()[..., None, None, None] * noisy_latents
+    )
+
+    true_log_variance_clipped = posterior_log_variance_clipped.to(device=timesteps.device)[timesteps].float()[
+        ..., None, None, None]
+
+    if noise_scheduler.variance_type == "learned":
+        model_log_variance = model_pred_var
+        #model_pred_var = th.exp(model_log_variance)
+    else:
+        min_log = true_log_variance_clipped
+        max_log = th.log(noise_scheduler.betas.to(device=timesteps.device)[timesteps].float()[..., None, None, None])
+        frac = (model_pred_var + 1) / 2
+        model_log_variance = frac * max_log + (1 - frac) * min_log
+        #model_pred_var = th.exp(model_log_variance)
+
+    sqrt_recip_alphas_cumprod = th.sqrt(1.0 / noise_scheduler.alphas_cumprod)
+    sqrt_recipm1_alphas_cumprod = th.sqrt(1.0 / noise_scheduler.alphas_cumprod - 1)
+
+    pred_xstart = (sqrt_recip_alphas_cumprod.to(device=timesteps.device)[timesteps].float()[
+                       ..., None, None, None] * noisy_latents
+                   - sqrt_recipm1_alphas_cumprod.to(device=timesteps.device)[timesteps].float()[
+                       ..., None, None, None] * model_pred_mean)
+
+    model_mean = (
+            posterior_mean_coef1.to(device=timesteps.device)[timesteps].float()[..., None, None, None] * pred_xstart
+            + posterior_mean_coef2.to(device=timesteps.device)[timesteps].float()[..., None, None, None] * noisy_latents
+    )
+
+    # model_mean = out["mean"] model_log_variance = out["log_variance"]
+    kl = normal_kl(
+        true_mean, true_log_variance_clipped, model_mean, model_log_variance
+    )
+    kl = kl.mean() / np.log(2.0)
+
+    decoder_nll = -discretized_gaussian_log_likelihood(
+        latents, means=model_mean, log_scales=0.5 * model_log_variance
+    )
+    assert decoder_nll.shape == latents.shape
+    decoder_nll = decoder_nll.mean() / np.log(2.0)
+
+    # At the first timestep return the decoder NLL,
+    # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+    kl_loss = th.where((timesteps == 0), decoder_nll, kl).mean()
+    return kl_loss
+
+def get_variance(noise_scheduler):
+    alphas_cumprod_prev = th.cat([th.tensor([1.0]), noise_scheduler.alphas_cumprod[:-1]])
+
+    posterior_mean_coef1 = (
+            noise_scheduler.betas * th.sqrt(alphas_cumprod_prev) / (1.0 - noise_scheduler.alphas_cumprod)
+    )
+
+    posterior_mean_coef2 = (
+            (1.0 - alphas_cumprod_prev)
+            * th.sqrt(noise_scheduler.alphas)
+            / (1.0 - noise_scheduler.alphas_cumprod)
+    )
+
+    posterior_variance = (
+            noise_scheduler.betas * (1.0 - alphas_cumprod_prev) / (1.0 - noise_scheduler.alphas_cumprod)
+    )
+    posterior_log_variance_clipped = th.log(
+        th.cat([posterior_variance[1][..., None], posterior_variance[1:]])
+    )
+    #res = posterior_log_variance_clipped.to(device=timesteps.device)[timesteps].float()
+    return posterior_mean_coef1, posterior_mean_coef2, posterior_log_variance_clipped #res[..., None, None, None]

diffusion_module/utils/noise_sampler.py

@@ -0,0 +1,16 @@
+from sklearn.mixture import GaussianMixture
+
+def get_noise_sampler(sample_type='gau'):
+    if sample_type == 'gau':
+        sampler = lambda latnt_sz: torch.randn_like(latnt_sz)
+    elif sample_type == 'gau_offset':
+        sampler = lambda latnt_sz: torch.randn_like(latnt_sz) + (torch.randn_like(latnt_sz))
+        ...
+    elif sample_type == 'gmm':
+        ...
+    else:
+        ...
+    return 
+
+if __name__ == "__main__":
+    ...

diffusion_module/utils/scheduler_factory.py

@@ -0,0 +1,300 @@
+from functools import partial
+from diffusers import DDPMScheduler, DPMSolverMultistepScheduler, UniPCMultistepScheduler, DPMSolverSinglestepScheduler
+from diffusers.pipeline_utils import DiffusionPipeline
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from typing import List, Optional, Tuple, Union
+import numpy as np
+from diffusers.schedulers.scheduling_utils import SchedulerOutput
+from diffusers.schedulers.scheduling_ddpm import DDPMSchedulerOutput
+from diffusers.utils import randn_tensor, BaseOutput
+
+
+### Testing the DDPM Scheduler for Variant 
+class ModifiedDDPMScheduler(DDPMScheduler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+    
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        generator=None,
+        return_dict: bool = True,
+    ) -> Union[DDPMSchedulerOutput, Tuple]:
+        """
+        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
+            timestep (`int`): current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+            generator: random number generator.
+            return_dict (`bool`): option for returning tuple rather than DDPMSchedulerOutput class
+
+        Returns:
+            [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] or `tuple`:
+            [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+
+        """
+        t = timestep
+
+        prev_t = self.previous_timestep(t)
+
+        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
+            print("Conidtion is trigger")
+       
+            model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
+            # [2,3, 64, 128]
+        else:
+            predicted_variance = None
+
+        # 1. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[t]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+        current_alpha_t = alpha_prod_t / alpha_prod_t_prev
+        current_beta_t = 1 - current_alpha_t
+
+        # 2. compute predicted original sample from predicted noise also called
+        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+        if self.config.prediction_type == "epsilon":
+            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+        
+        elif self.config.prediction_type == "sample":
+            pred_original_sample = model_output
+        elif self.config.prediction_type == "v_prediction":
+            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+                " `v_prediction`  for the DDPMScheduler."
+            )
+
+        # 3. Clip or threshold "predicted x_0"
+        if self.config.thresholding:
+            pred_original_sample = self._threshold_sample(pred_original_sample)
+        elif self.config.clip_sample:
+            pred_original_sample = pred_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
+        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+        pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t
+        current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
+
+        # 5. Compute predicted previous sample µ_t
+        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
+
+        # 6. Add noise
+        variance = 0
+        if t > 0:
+            device = model_output.device
+            variance_noise = randn_tensor(
+                model_output.shape, generator=generator, device=device, dtype=model_output.dtype
+            )
+            if self.variance_type == "fixed_small_log":
+                variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise
+            
+            elif self.variance_type == "learned_range":
+                variance = self._get_variance(t, predicted_variance=predicted_variance)
+                variance = torch.exp(0.5 * variance) * variance_noise
+
+            else:
+                variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise
+        
+        pred_prev_sample = pred_prev_sample + variance
+        print(pred_prev_sample.shape)
+        if not return_dict:
+            return (pred_prev_sample,)
+
+        return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
+   
+
+class ModifiedUniPCScheduler(UniPCMultistepScheduler):
+    '''
+    This is the modification of UniPCMultistepScheduler, which is the same as UniPCMultistepScheduler except for the _get_variance function.
+    '''
+    def __init__(self, variance_type: str = "fixed_small", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.custom_timesteps = False
+        self.variance_type=variance_type
+        self.config.timestep_spacing="leading"
+    def previous_timestep(self, timestep):
+        if self.custom_timesteps:
+            index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0]
+            if index == self.timesteps.shape[0] - 1:
+                prev_t = torch.tensor(-1)
+            else:
+                prev_t = self.timesteps[index + 1]
+        else:
+            num_inference_steps = (
+                self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
+            )
+            prev_t = timestep - self.config.num_train_timesteps // num_inference_steps
+
+        return prev_t
+    
+    def _get_variance(self, t, predicted_variance=None, variance_type="learned_range"):
+        prev_t = self.previous_timestep(t)
+
+        alpha_prod_t = self.alphas_cumprod[t]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
+        current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev
+
+        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
+
+        variance = torch.clamp(variance, min=1e-20)
+
+        if variance_type is None:
+            variance_type = self.config.variance_type
+
+        if variance_type == "fixed_small":
+            variance = variance
+        elif variance_type == "fixed_small_log":
+            variance = torch.log(variance)
+            variance = torch.exp(0.5 * variance)
+        elif variance_type == "fixed_large":
+            variance = current_beta_t
+        elif variance_type == "fixed_large_log":
+            variance = torch.log(current_beta_t)
+        elif variance_type == "learned":
+            return predicted_variance
+        elif variance_type == "learned_range":
+            min_log = torch.log(variance)
+            max_log = torch.log(current_beta_t)
+            frac = (predicted_variance + 1) / 2
+            variance = frac * max_log + (1 - frac) * min_log
+
+        return variance
+
+    def step(self, model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, return_dict: bool = True) -> Union[SchedulerOutput, Tuple]:
+        
+        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
+            print("condition using predicted_variance is trigger")
+            model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
+        else:
+            predicted_variance = None
+
+        super_output = super().step(model_output, timestep, sample, return_dict=False)
+        prev_sample = super_output[0]
+        # breakpoint()
+        variance = 0
+        if timestep > 0:
+            device = model_output.device
+            variance_noise = randn_tensor(
+                model_output.shape, generator=None, device=device, dtype=model_output.dtype
+            )
+            if self.variance_type == "fixed_small_log":
+                variance = self._get_variance(timestep, predicted_variance=predicted_variance) * variance_noise
+            elif self.variance_type == "learned_range":
+                # breakpoint()
+                variance = self._get_variance(timestep, predicted_variance=predicted_variance)
+                variance = torch.exp(0.5 * variance) * variance_noise
+                # breakpoint()
+            else:
+                variance = (self._get_variance(timestep, predicted_variance=predicted_variance) ** 0.5) * variance_noise
+
+      
+        # breakpoint()
+        print("time step is ", timestep)
+        prev_sample = prev_sample  + variance
+
+        if not return_dict:
+            return (prev_sample,)
+        
+        return DDPMSchedulerOutput(prev_sample=prev_sample,pred_original_sample=prev_sample) 
+
+        #return SchedulerOutput(prev_sample=prev_sample)
+
+
+def build_proc(sch_cfg=None, _sch=None, **kwargs):
+    if kwargs:
+        return _sch(**kwargs)
+
+    type_str = str(type(sch_cfg))
+    if 'dict' in type_str:
+        return _sch.from_config(**sch_cfg)
+    return _sch.from_config(sch_cfg, subfolder="scheduler")
+
+scheduler_factory = {
+    'UniPC' : partial(build_proc, _sch=UniPCMultistepScheduler),
+    'modifiedUniPC' : partial(build_proc, _sch=ModifiedUniPCScheduler),
+    # DPM family
+    'DDPM' : partial(build_proc, _sch=DDPMScheduler),
+    'DPMSolver' : partial(build_proc, _sch=DPMSolverMultistepScheduler, algorithm_type='dpmsolver'),
+    'DPMSolver++' : partial(build_proc, _sch=DPMSolverMultistepScheduler),
+    'DPMSolverSingleStep' : partial(build_proc, _sch=DPMSolverSinglestepScheduler)
+
+}
+
+def scheduler_setup(pipe : DiffusionPipeline = None, scheduler_type : str = 'UniPC', from_config=None, **kwargs):
+    if not isinstance(pipe, DiffusionPipeline):
+        raise TypeError(f'pipe should be DiffusionPipeline, but given {type(pipe)}\n')
+
+    sch_cfg = from_config if from_config else pipe.scheduler.config  
+    #sch_cfg = diffusers.configuration_utils.FrozenDict({**sch_cfg, 'solver_order':3})  
+    #pipe.scheduler = scheduler_factory[scheduler_type](**kwargs) if kwargs \
+    #                    else scheduler_factory[scheduler_type](sch_cfg)
+    
+    # pipe.scheduler = DPMSolverSinglestepScheduler()
+    # #pipe.scheduler = DDPMScheduler(beta_schedule="linear", variance_type="learned_range")
+    # print(pipe.scheduler)
+    print("Scheduler type in Scheduler_factory.py is Hard-coded to modifyUniPC, Please change it back to AutoDetect functionality if you want to change scheudler")
+    pipe.scheduler = ModifiedUniPCScheduler(variance_type="learned_range", )
+    # pipe.scheduler = ModifiedDDPMScheduler(beta_schedule="linear", variance_type="learned_range")
+    
+    #pipe.scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
+    #pipe.scheduler._get_variance = _get_variance
+    return pipe
+
+# unittest of scheduler..
+if __name__ == "__main__":
+    def ld_mod():   
+        noise_scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
+        vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to("cuda").to(torch.float16)
+        unet = SDMUNet2DModel.from_pretrained("/data/harry/Data_generation/diffusers-main/examples/VAESDM/LDM-sdm-model/checkpoint-46000", subfolder="unet").to("cuda").to(torch.float16)
+        return noise_scheduler, vae, unet
+
+    from Pipline import SDMLDMPipeline
+    from diffusers import StableDiffusionPipeline
+    import torch
+
+    path = "CompVis/stable-diffusion-v1-4"
+    pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+    
+    # change scheduler 
+    # customized args : once you customized, customize forever ~ no from_config
+    #pipe = scheduler_setup(pipe, 'DPMSolver++', thresholding=True)
+    # from_config
+    pipe = scheduler_setup(pipe, 'DPMSolverSingleStep')
+
+    pipe = pipe.to("cuda")
+    prompt = "a highly realistic photo of green turtle"
+    generator = torch.manual_seed(0)
+    # only 15 steps are needed for good results => 2-4 seconds on GPU
+    image = pipe(prompt, generator=generator, num_inference_steps=15).images[0]
+    # save image
+    image.save("turtle.png")
+
+    '''
+    # load & wrap submodules into pipe-API
+    noise_scheduler, vae, unet = ld_mod()
+    pipe = SDMLDMPipeline(
+        unet=unet,
+        vqvae=vae,
+        scheduler=noise_scheduler,
+        torch_dtype=torch.float16
+    )
+
+    # change scheduler 
+    pipe = scheduler_setup(pipe, 'DPMSolverSingleStep')
+    pipe = pipe.to("cuda")
+    '''

evolution.py

@@ -0,0 +1,102 @@
+import random
+import math
+import numpy as np
+from PIL import Image
+from skimage.draw import line
+from skimage import morphology
+import cv2
+
+def line_crosses_cracks(start, end, img):
+    rr, cc = line(start[0], start[1], end[0], end[1])
+    # Exclude the starting point from the line coordinates
+    if len(rr) > 1 and len(cc) > 1:
+        return np.any(img[rr[1:], cc[1:]] == 255)
+    return False
+
+def random_walk(img_array, k=8, m=0.1, min_steps=50, max_steps=200, length=2, degree_range=30, seed=None):
+    
+    if seed is not None:
+        random.seed(seed)
+        np.random.seed(seed)
+
+    
+    img_array = cv2.ximgproc.thinning(img_array)
+
+    rows, cols = img_array.shape
+    # Find all white pixels (existing cracks)
+    white_pixels = np.column_stack(np.where(img_array == 255))
+    original_crack_count = len(white_pixels)  # Count of original crack pixels
+
+    # Select k random starting points from the white pixels
+    if white_pixels.size == 0:
+        raise ValueError("No initial crack pixels found in the image.")
+    if k > len(white_pixels):
+        raise ValueError("k is greater than the number of existing crack pixels.")
+    initial_points = white_pixels[random.sample(range(len(white_pixels)), k)]
+
+    # Initialize step count for each initial point with a random value between min_steps and max_steps
+    step_counts = {i: random.randint(min_steps, max_steps) for i in range(k)}
+    # Initialize main direction for each initial point (0 to 360 degrees)
+    main_angles = {i: random.uniform(0, 360) for i in range(k)}
+
+    grown_crack_count = 0  # Count of newly grown crack pixels
+
+    # Start the random walk for each initial point
+    for idx, point in enumerate(initial_points):
+        current_pos = tuple(point)
+        current_steps = 0
+        while current_steps < step_counts[idx]:
+            # Check the crack ratio
+            current_ratio = np.sum(img_array == 255) / (rows * cols)
+            if current_ratio >= m:
+                return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}
+
+            # Generate a random direction within the fan-shaped area around the main angle
+            main_angle = main_angles[idx]
+            angle = math.radians(main_angle + random.uniform(-degree_range, degree_range))
+            
+            # Determine the next position with the specified length
+            delta_row = length * math.sin(angle)
+            delta_col = length * math.cos(angle)
+            next_pos = (int(current_pos[0] + delta_row), int(current_pos[1] + delta_col))
+            
+            # Check if the line from the current to the next position crosses existing cracks
+            if 0 <= next_pos[0] < rows and 0 <= next_pos[1] < cols and not line_crosses_cracks(current_pos, next_pos, img_array):
+                # Draw a line from the current position to the next position
+                rr, cc = line(current_pos[0], current_pos[1], next_pos[0], next_pos[1])
+                img_array[rr, cc] = 255  # Set the pixels along the line to white
+                grown_crack_count += len(rr)  # Update the count of grown crack pixels
+                current_pos = next_pos
+                current_steps += 1
+            else:
+                # If the line crosses existing cracks or the next position is outside the boundaries, stop the walk for this point
+                break
+
+    return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}
+
+# The rest of the test code remains the same.
+# You can use this function in your test code to generate the image and get the counts.
+
+
+# test code 
+if __name__ == "__main__":
+    # Updated parameters
+    k = 8  # Number of initial white pixels to start the random walk
+    m = 0.1  # Maximum ratio of crack pixels
+    min_steps = 50
+    max_steps = 200
+    img_path = '/data/leiqin/diffusion/huggingface_diffusers/crack_label_creator/random_walk/thindata_256/2.png'
+    img = Image.open(img_path)
+    img_array = np.array(img)
+    length = 2
+
+    # Perform the modified random walk
+    result_img_array_mod, pixels_dict = random_walk(img_array.copy(), k, m, min_steps, max_steps, length)
+
+    # Convert the result to an image
+    result_img_mod = Image.fromarray(result_img_array_mod.astype('uint8'))
+
+    # Save the resulting image
+    result_img_path_mod = 'resutls.png'
+    result_img_mod.save(result_img_path_mod)
+    print(pixels_dict)

generate.py

@@ -0,0 +1,106 @@
+from diffusers.schedulers import UniPCMultistepScheduler    
+from diffusers import AutoencoderKL
+from diffusion_module.unet import UNetModel
+import torch
+from diffusion_module.utils.LSDMPipeline_expandDataset import SDMLDMPipeline
+from accelerate import Accelerator
+from evolution import random_walk
+import cv2
+import numpy as np 
+
+def mask2onehot(data, num_classes):
+    # move to GPU and change data types
+    data = data.to(dtype=torch.int64)
+
+    # create one-hot label map
+    label_map = data
+    bs, _, h, w = label_map.size()
+    input_label = torch.FloatTensor(bs, num_classes, h, w).zero_().to(data.device)
+    input_semantics = input_label.scatter_(1, label_map, 1.0)
+
+    return input_semantics
+
+def generate(img, pretrain_weight,seed=None):
+
+    noise_scheduler = UniPCMultistepScheduler()
+    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
+    latent_size  = (64, 64)
+    unet = UNetModel(
+        image_size = latent_size,
+        in_channels=vae.config.latent_channels,
+        model_channels=256,
+        out_channels=vae.config.latent_channels,
+        num_res_blocks=2,
+        attention_resolutions=(2, 4, 8),
+        dropout=0,
+        channel_mult=(1, 2, 3, 4),
+        num_heads=8,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=True,
+        resblock_updown=True,
+        use_new_attention_order=False,
+        num_classes=151,
+        mask_emb="resize",
+        use_checkpoint=True,
+        SPADE_type="spade",
+    )
+    
+
+    unet = unet.from_pretrained(pretrain_weight)
+    device = 'cpu'
+    if device != 'cpu':
+        mixed_precision = "fp16"
+    else:
+        mixed_precision = "no"
+    
+    
+    accelerator = Accelerator(
+        mixed_precision=mixed_precision,
+        cpu= True if device is 'cpu' else False
+    )
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    
+    unet,vae = accelerator.prepare(unet, vae)
+    vae.to(device=accelerator.device, dtype=weight_dtype)
+    pipeline = SDMLDMPipeline(
+        vae=accelerator.unwrap_model(vae),
+        unet=accelerator.unwrap_model(unet),
+        scheduler=noise_scheduler,
+        torch_dtype=weight_dtype,
+        resolution_type="crack"
+    )
+    """
+    if accelerator.device != 'cpu':
+        pipeline.enable_xformers_memory_efficient_attention()
+    """
+    pipeline = pipeline.to(accelerator.device)
+    pipeline.set_progress_bar_config(disable=False)
+
+    if seed is None:
+        generator = None
+    else:
+        generator = torch.Generator(device=accelerator.device).manual_seed(seed)
+    
+    resized_s = cv2.resize(img, (64, 64), interpolation=cv2.INTER_AREA)
+    # 灰度图放大到255
+    _, binary_s = cv2.threshold(resized_s, 1, 255, cv2.THRESH_BINARY)
+    # 转换为0,1
+    tensor_s = torch.from_numpy(binary_s / 255)
+    # h,w -> 1,1,h,w
+    tensor_s = tensor_s.unsqueeze(0).unsqueeze(0)
+    onehot_skeletons=[]
+    onehot_s = mask2onehot(tensor_s, 151)
+    onehot_skeletons.append(onehot_s)
+
+    onehot_skeletons = torch.stack(onehot_skeletons, dim=1).squeeze(0)
+    onehot_skeletons = onehot_skeletons.to(dtype=weight_dtype,device=accelerator.device)
+
+    images = pipeline(onehot_skeletons, generator=generator,batch_size = 1,
+                    num_inference_steps=20, s=1.5,
+                    num_evolution_per_mask=1).images
+    
+    return images

requirements.txt

@@ -0,0 +1,137 @@
+accelerate==0.25.0
+aiofiles==23.2.1
+aiohttp==3.9.1
+aiosignal==1.3.1
+albumentations==1.3.1
+altair==5.2.0
+annotated-types==0.6.0
+antlr4-python3-runtime==4.9.3
+anyio==4.2.0
+appdirs==1.4.4
+async-timeout==4.0.3
+attrs==23.2.0
+blobfile==2.1.1
+certifi==2023.11.17
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+contourpy==1.2.0
+cycler==0.12.1
+datasets==2.16.1
+diffusers==0.16.1
+dill==0.3.7
+docker-pycreds==0.4.0
+einops==0.7.0
+exceptiongroup==1.2.0
+fastapi==0.109.2
+ffmpy==0.3.1
+filelock==3.13.1
+fonttools==4.47.0
+frozenlist==1.4.1
+fsspec==2023.10.0
+gitdb==4.0.11
+GitPython==3.1.41
+gradio==4.16.0
+gradio_client==0.8.1
+h11==0.14.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.20.1
+idna==3.6
+imageio==2.33.1
+importlib-metadata==7.0.1
+importlib-resources==6.1.1
+Jinja2==3.1.2
+joblib==1.3.2
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+lazy_loader==0.3
+lightning-utilities==0.10.0
+lxml==4.9.4
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.2
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+multiprocess==0.70.15
+networkx==3.2.1
+numpy==1.26.2
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+opencv-contrib-python==4.9.0.80
+opencv-python-headless==4.9.0.80
+orjson==3.9.13
+packaging==23.2
+pandas==2.1.4
+pillow==10.2.0
+protobuf==4.25.2
+psutil==5.9.7
+pyarrow==14.0.2
+pyarrow-hotfix==0.6
+pycryptodomex==3.20.0
+pydantic==2.6.0
+pydantic_core==2.16.1
+pydub==0.25.1
+Pygments==2.17.2
+pyparsing==3.1.1
+python-dateutil==2.8.2
+python-multipart==0.0.7
+pytorch-lightning==2.1.3
+pytz==2023.3.post1
+PyYAML==6.0.1
+qudida==0.0.4
+referencing==0.33.0
+regex==2023.12.25
+requests==2.31.0
+rich==13.7.0
+rpds-py==0.17.1
+ruff==0.2.0
+safetensors==0.4.1
+scikit-image==0.22.0
+scikit-learn==1.3.2
+scipy==1.11.4
+semantic-version==2.10.0
+sentry-sdk==1.39.2
+setproctitle==1.3.3
+shellingham==1.5.4
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.0
+starlette==0.36.3
+sympy==1.12
+threadpoolctl==3.2.0
+tifffile==2023.12.9
+tokenizers==0.15.0
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.1.2
+torchaudio==2.1.2
+torchmetrics==1.2.1
+torchvision==0.16.2
+tqdm==4.66.1
+transformers==4.36.2
+triton==2.1.0
+typer==0.9.0
+typing_extensions==4.9.0
+tzdata==2023.4
+urllib3==2.1.0
+uvicorn==0.27.0.post1
+wandb==0.16.2
+websockets==11.0.3
+xformers==0.0.23.post1
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.17.0