xl version

NoiseTransformer.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+
+from torch.nn import functional as F
+from timm import create_model
+
+
+__all__ = ['NoiseTransformer']
+
+class NoiseTransformer(nn.Module):
+    def __init__(self, resolution=(128,96)):
+        super().__init__()
+        self.upsample = lambda x: F.interpolate(x, [224,224])
+        self.downsample = lambda x: F.interpolate(x, [resolution[0],resolution[1]])
+        self.upconv = nn.Conv2d(7,4,(1,1),(1,1),(0,0))
+        self.downconv = nn.Conv2d(4,3,(1,1),(1,1),(0,0))
+        # self.upconv = nn.Conv2d(7,4,(1,1),(1,1),(0,0))
+        self.swin = create_model("swin_tiny_patch4_window7_224",pretrained=True)
+
+
+    def forward(self, x, residual=False):
+        if residual:
+            x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x))))) + x
+        else:
+            x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x)))))
+
+        return x

README.md

@@ -1,5 +1,5 @@
 ---
-title: Hyperparameters Are All You Need Xl Version
+title: Hyperparameters-are-all-you-need-xl-version-improved-implementation
 emoji: 🖼
 colorFrom: purple
 colorTo: red
@@ -8,7 +8,21 @@ sdk_version: 5.44.0
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: training-free algorithm for few step inference of sdxl
+short_description: training-free few step diffusion ODE solver
+header: default
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+**Abstract:** The diffusion model is a state-of-the-art generative model that generates an image by applying a neural network iteratively. Moreover, this generation process is regarded as an algorithm solving an ordinary differential equation or a stochastic differential equation. Based on the analysis of the truncation error of the diffusion ODE and SDE, our study proposes a training-free algorithm that generates high-quality 512 x 512 and 1024 x 1024 images in eight steps, with flexible guidance scales. To the best of my knowledge, our algorithm is the first one that samples a 1024 x 1024 resolution image in 8 steps with an FID performance comparable to that of the latest distillation model, but without additional training. Meanwhile, our algorithm can also generate a 512 x 512 image in 8 steps, and its FID performance is better than the inference result using state-of-the-art ODE solver DPM++ 2m in 20 steps. We validate our eight-step image generation algorithm using the COCO 2014, COCO 2017, and LAION datasets. And our best FID performance is 15.7, 22.35, and 17.52. While the FID performance of DPM++2m is 17.3, 23.75, and 17.33. Further, it also outperforms the state-of-the-art AMED-plugin solver, whose FID performance is 19.07, 25.50, and 18.06. We also apply the algorithm in five-step inference without additional training, for which the best FID performance in the datasets mentioned above is 19.18, 23.24, and 19.61, respectively, and is comparable to the performance of the state-of-the-art AMED Pulgin solver in eight steps, SDXL-turbo in four steps, and the state-of-the-art diffusion distillation model Flash Diffusion in five steps. We also validate our algorithm in synthesizing 1024 * 1024 images within 6 steps, whose FID performance only has a limited distance to the latest distillation algorithm. 
+
+This is a demo is a simplified version of the approach described in the paper, ["Hyperparameters are all you need: Using five-step inference for an original diffusion model to generate images comparable to the latest distillation model"](https://arxiv.org/abs/2510.02390)
+
+```
+@misc{hyper,
+      title={Hyperparameters are all you need: Using five-step inference for an original diffusion model to generate images comparable to the latest distillation model}, 
+      author={Zilai Li},
+      year={2025},
+      eprint={2510.02390},
+      archivePrefix={arXiv},
+      primaryClass={eess.IV}
+}
+```

SVDNoiseUnet.py

@@ -0,0 +1,430 @@
+import torch
+import torch.nn as nn
+import einops
+
+from torch.nn import functional as F
+from torch.jit import Final
+from timm.layers import use_fused_attn
+from timm.models.layers import PatchEmbed, Mlp, DropPath, trunc_normal_, lecun_normal_, get_act_layer
+from abc import abstractmethod
+from NoiseTransformer import NoiseTransformer
+from einops import rearrange
+__all__ = ['SVDNoiseUnet', 'SVDNoiseUnet_Concise']
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SVDNoiseUnet(nn.Module):
+    def __init__(self, in_channels=4, out_channels=4, resolution=(128,96)): # resolution = size // 8
+        super(SVDNoiseUnet, self).__init__()
+
+        _in_1 = int(resolution[0] * in_channels // 2)
+        _out_1 = int(resolution[0] * out_channels // 2)
+        
+        _in_2 = int(resolution[1] * in_channels // 2)
+        _out_2 = int(resolution[1] * out_channels // 2)
+        self.mlp1 = nn.Sequential(
+            nn.Linear(_in_1, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out_1),
+        )
+        self.mlp2 = nn.Sequential(
+            nn.Linear(_in_2, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out_2),
+        )
+
+        self.mlp3 = nn.Sequential(
+            nn.Linear(_in_2, _out_2),
+        )
+
+        self.attention = Attention(_out_2)
+
+        self.bn = nn.BatchNorm1d(256)
+        self.bn2 = nn.BatchNorm1d(192)
+
+        self.mlp4 =  nn.Sequential(
+            nn.Linear(_out_2, 1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, _out_2),
+        )
+        self.ffn = nn.Sequential(
+            nn.Linear(256, 384),  # Expand
+            nn.ReLU(inplace=True),
+            nn.Linear(384, 192)   # Reduce to target size
+        )
+        self.ffn2 = nn.Sequential(
+            nn.Linear(256, 384),  # Expand
+            nn.ReLU(inplace=True),
+            nn.Linear(384, 192)   # Reduce to target size
+        )
+        # self.adaptive_pool = nn.AdaptiveAvgPool2d((256, 192))
+
+    def forward(self, x, residual=False):
+        b, c, h, w = x.shape
+        x = einops.rearrange(x, "b (a c)h w ->b (a h)(c w)", a=2,c=2) # x -> [1, 256, 256]
+        U, s, V = torch.linalg.svd(x) # U->[b 256 256], s-> [b 256], V->[b 256 256]
+        U_T = U.permute(0, 2, 1)
+        U_out = self.ffn(self.mlp1(U_T))
+        U_out = self.bn(U_out)
+        U_out = U_out.transpose(1, 2)
+        U_out = self.ffn2(U_out)  # [b, 256, 256] -> [b, 256, 192]
+        U_out = self.bn2(U_out)
+        U_out = U_out.transpose(1, 2)
+        # U_out = self.bn(U_out)
+        V_out = self.mlp2(V)
+        s_out = self.mlp3(s).unsqueeze(1)  # s -> [b, 1, 256]  => [b, 256, 256]
+        out = U_out + V_out + s_out
+        # print(out.size())
+        out = out.squeeze(1)
+        out = self.attention(out).mean(1)
+        out = self.mlp4(out) + s
+        diagonal_out = torch.diag_embed(out)
+        padded_diag = F.pad(diagonal_out, (0, 0, 0, 64), mode='constant', value=0)  # Shape: [b, 1, 256, 192]
+        pred = U @ padded_diag @ V
+        return einops.rearrange(pred, "b (a h)(c w) -> b (a c) h w", a=2,c=2)
+    
+class SVDNoiseUnet64(nn.Module):
+    def __init__(self, in_channels=4, out_channels=4, resolution=64): # resolution = size // 8
+        super(SVDNoiseUnet64, self).__init__()
+
+        _in = int(resolution * in_channels // 2)
+        _out = int(resolution * out_channels // 2)
+        self.mlp1 = nn.Sequential(
+            nn.Linear(_in, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out),
+        )
+        self.mlp2 = nn.Sequential(
+            nn.Linear(_in, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out),
+        )
+
+        self.mlp3 = nn.Sequential(
+            nn.Linear(_in, _out),
+        )
+
+        self.attention = Attention(_out)
+
+        self.bn = nn.BatchNorm2d(_out)
+
+        self.mlp4 =  nn.Sequential(
+            nn.Linear(_out, 1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, _out),
+        )
+
+    def forward(self, x, residual=False):
+        b, c, h, w = x.shape
+        x = einops.rearrange(x, "b (a c)h w ->b (a h)(c w)", a=2,c=2) # x -> [1, 256, 256]
+        U, s, V = torch.linalg.svd(x) # U->[b 256 256], s-> [b 256], V->[b 256 256]
+        U_T = U.permute(0, 2, 1)
+        out = self.mlp1(U_T) + self.mlp2(V) + self.mlp3(s).unsqueeze(1) # s -> [b, 1, 256]  => [b, 256, 256]
+        out = self.attention(out).mean(1)
+        out = self.mlp4(out) + s
+        pred = U @ torch.diag_embed(out) @ V
+        return einops.rearrange(pred, "b (a h)(c w) -> b (a c) h w", a=2,c=2)
+    
+
+
+class SVDNoiseUnet128(nn.Module):
+    def __init__(self, in_channels=4, out_channels=4, resolution=128): # resolution = size // 8
+        super(SVDNoiseUnet128, self).__init__()
+
+        _in = int(resolution * in_channels // 2)
+        _out = int(resolution * out_channels // 2)
+        self.mlp1 = nn.Sequential(
+            nn.Linear(_in, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out),
+        )
+        self.mlp2 = nn.Sequential(
+            nn.Linear(_in, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, _out),
+        )
+
+        self.mlp3 = nn.Sequential(
+            nn.Linear(_in, _out),
+        )
+
+        self.attention = Attention(_out)
+
+        self.bn = nn.BatchNorm2d(_out)
+
+        self.mlp4 =  nn.Sequential(
+            nn.Linear(_out, 1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, _out),
+        )
+
+    def forward(self, x, residual=False):
+        b, c, h, w = x.shape
+        x = einops.rearrange(x, "b (a c)h w ->b (a h)(c w)", a=2,c=2) # x -> [1, 256, 256]
+        U, s, V = torch.linalg.svd(x) # U->[b 256 256], s-> [b 256], V->[b 256 256]
+        U_T = U.permute(0, 2, 1)
+        out = self.mlp1(U_T) + self.mlp2(V) + self.mlp3(s).unsqueeze(1) # s -> [b, 1, 256]  => [b, 256, 256]
+        out = self.attention(out).mean(1)
+        out = self.mlp4(out) + s
+        pred = U @ torch.diag_embed(out) @ V
+        return einops.rearrange(pred, "b (a h)(c w) -> b (a c) h w", a=2,c=2)
+    
+
+
+class SVDNoiseUnet_Concise(nn.Module):
+    def __init__(self, in_channels=4, out_channels=4, resolution=64):
+        super(SVDNoiseUnet_Concise, self).__init__()
+        
+
+from diffusers.models.normalization import AdaGroupNorm
+
+class NPNet(nn.Module):
+      def __init__(self, model_id, pretrained_path=' ', device='cuda') -> None:
+            super(NPNet, self).__init__()
+
+            assert model_id in ['SD1.5', 'DreamShaper', 'DiT']
+
+            self.model_id = model_id
+            self.device = device
+            self.pretrained_path = pretrained_path
+
+            (
+                  self.unet_svd, 
+                  self.unet_embedding, 
+                  self.text_embedding, 
+                  self._alpha, 
+                  self._beta
+             ) = self.get_model()
+      def save_model(self, save_path: str):
+            """
+            Save this NPNet so that get_model() can later reload it.
+            """
+            torch.save({
+                  "unet_svd":        self.unet_svd.state_dict(),
+                  "unet_embedding":  self.unet_embedding.state_dict(),
+                  "embeeding":       self.text_embedding.state_dict(),  # matches get_model’s key
+                  "alpha":           self._alpha,
+                  "beta":            self._beta,
+            }, save_path)
+            print(f"NPNet saved to {save_path}")
+      def get_model(self):
+
+            unet_embedding = NoiseTransformer(resolution=(128,96)).to(self.device).to(torch.float32)
+            unet_svd = SVDNoiseUnet(resolution=(128,96)).to(self.device).to(torch.float32)
+
+            if self.model_id == 'DiT':
+                  text_embedding = AdaGroupNorm(768 * 77, 4, 1, eps=1e-6).to(self.device).to(torch.float32)
+            else:
+                  text_embedding = AdaGroupNorm(768 * 77, 4, 1, eps=1e-6).to(self.device).to(torch.float32) 
+
+            # initialize random _alpha and _beta when no checkpoint is provided
+            _alpha = torch.randn(1, device=self.device)
+            _beta = torch.randn(1, device=self.device)
+
+            if '.pth' in self.pretrained_path:
+                  gloden_unet = torch.load(self.pretrained_path)
+                  unet_svd.load_state_dict(gloden_unet["unet_svd"],strict=True)
+                  unet_embedding.load_state_dict(gloden_unet["unet_embedding"],strict=True)
+                  text_embedding.load_state_dict(gloden_unet["embeeding"],strict=True)
+                  _alpha = gloden_unet["alpha"]
+                  _beta = gloden_unet["beta"]
+
+                  print("Load Successfully!")
+
+                  return unet_svd, unet_embedding, text_embedding, _alpha, _beta
+            
+            else:
+                  return unet_svd, unet_embedding, text_embedding, _alpha, _beta
+            
+
+      def forward(self, initial_noise, prompt_embeds):
+
+            prompt_embeds = prompt_embeds.float().view(prompt_embeds.shape[0], -1)
+            text_emb = self.text_embedding(initial_noise.float(), prompt_embeds)
+
+            encoder_hidden_states_svd = initial_noise
+            encoder_hidden_states_embedding = initial_noise + text_emb
+
+            golden_embedding = self.unet_embedding(encoder_hidden_states_embedding.float())
+
+            golden_noise = self.unet_svd(encoder_hidden_states_svd.float()) + (
+                        2 * torch.sigmoid(self._alpha) - 1) * text_emb + self._beta * golden_embedding
+
+            return golden_noise
+
+
+class NPNet64(nn.Module):
+      def __init__(self, model_id, pretrained_path=' ', device='cuda') -> None:
+            super(NPNet64, self).__init__()
+            self.model_id = model_id
+            self.device = device
+            self.pretrained_path = pretrained_path
+
+            (
+                  self.unet_svd, 
+                  self.unet_embedding, 
+                  self.text_embedding, 
+                  self._alpha, 
+                  self._beta
+             ) = self.get_model()
+
+      def save_model(self, save_path: str):
+            """
+            Save this NPNet so that get_model() can later reload it.
+            """
+            torch.save({
+                  "unet_svd":        self.unet_svd.state_dict(),
+                  "unet_embedding":  self.unet_embedding.state_dict(),
+                  "embeeding":       self.text_embedding.state_dict(),  # matches get_model’s key
+                  "alpha":           self._alpha,
+                  "beta":            self._beta,
+            }, save_path)
+            print(f"NPNet saved to {save_path}")
+        
+      def get_model(self):
+
+            unet_embedding = NoiseTransformer(resolution=(64,64)).to(self.device).to(torch.float32)
+            unet_svd = SVDNoiseUnet64(resolution=64).to(self.device).to(torch.float32)
+            _alpha = torch.randn(1, device=self.device)
+            _beta = torch.randn(1, device=self.device)
+
+            text_embedding = AdaGroupNorm(768 * 77, 4, 1, eps=1e-6).to(self.device).to(torch.float32) 
+
+            
+            if '.pth' in self.pretrained_path:
+                  gloden_unet = torch.load(self.pretrained_path)
+                  unet_svd.load_state_dict(gloden_unet["unet_svd"])
+                  unet_embedding.load_state_dict(gloden_unet["unet_embedding"])
+                  text_embedding.load_state_dict(gloden_unet["embeeding"])
+                  _alpha = gloden_unet["alpha"]
+                  _beta = gloden_unet["beta"]
+
+                  print("Load Successfully!")
+
+            return unet_svd, unet_embedding, text_embedding, _alpha, _beta
+            
+
+      def forward(self, initial_noise, prompt_embeds):
+
+            prompt_embeds = prompt_embeds.float().view(prompt_embeds.shape[0], -1)
+            text_emb = self.text_embedding(initial_noise.float(), prompt_embeds)
+
+            encoder_hidden_states_svd = initial_noise
+            encoder_hidden_states_embedding = initial_noise + text_emb
+
+            golden_embedding = self.unet_embedding(encoder_hidden_states_embedding.float())
+
+            golden_noise = self.unet_svd(encoder_hidden_states_svd.float()) + (
+                        2 * torch.sigmoid(self._alpha) - 1) * text_emb + self._beta * golden_embedding
+
+            return golden_noise
+
+class NPNet128(nn.Module):
+      def __init__(self, model_id, pretrained_path=True, device='cuda') -> None:
+            super(NPNet128, self).__init__()
+
+            assert model_id in ['SDXL', 'DreamShaper', 'DiT']
+
+            self.model_id = model_id
+            self.device = device
+            self.pretrained_path = pretrained_path
+
+            (
+                  self.unet_svd, 
+                  self.unet_embedding, 
+                  self.text_embedding, 
+                  self._alpha, 
+                  self._beta
+             ) = self.get_model()
+
+      def get_model(self):
+
+            unet_embedding = NoiseTransformer(resolution=(128,128)).to(self.device).to(torch.float32)
+            unet_svd = SVDNoiseUnet128(resolution=128).to(self.device).to(torch.float32)
+
+            if self.model_id == 'DiT':
+                  text_embedding = AdaGroupNorm(1024 * 77, 4, 1, eps=1e-6).to(self.device).to(torch.float32)
+            else:
+                  text_embedding = AdaGroupNorm(2048 * 77, 4, 1, eps=1e-6).to(self.device).to(torch.float32) 
+
+            
+            if '.pth' in self.pretrained_path:
+                  gloden_unet = torch.load(self.pretrained_path)
+                  unet_svd.load_state_dict(gloden_unet["unet_svd"])
+                  unet_embedding.load_state_dict(gloden_unet["unet_embedding"])
+                  text_embedding.load_state_dict(gloden_unet["embeeding"])
+                  _alpha = gloden_unet["alpha"]
+                  _beta = gloden_unet["beta"]
+
+                  print("Load Successfully!")
+
+                  return unet_svd, unet_embedding, text_embedding, _alpha, _beta
+            
+            else:
+                  assert ("No Pretrained Weights Found!")
+            
+
+      def forward(self, initial_noise, prompt_embeds):
+
+            prompt_embeds = prompt_embeds.float().view(prompt_embeds.shape[0], -1)
+            text_emb = self.text_embedding(initial_noise.float(), prompt_embeds)
+
+            encoder_hidden_states_svd = initial_noise
+            encoder_hidden_states_embedding = initial_noise + text_emb
+
+            golden_embedding = self.unet_embedding(encoder_hidden_states_embedding.float())
+
+            golden_noise = self.unet_svd(encoder_hidden_states_svd.float()) + (
+                        2 * torch.sigmoid(self._alpha) - 1) * text_emb + self._beta * golden_embedding
+
+            return golden_noise
+      

app.py

@@ -1,60 +1,966 @@
 import gradio as gr
 import numpy as np
 import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+import json
+import spaces #[uncomment to use ZeroGPU]
+from diffusers import (
+    AutoencoderKL,
+    StableDiffusionXLPipeline,
+)
+from huggingface_hub import login, hf_hub_download
+from PIL import Image
+# from huggingface_hub import login
+from SVDNoiseUnet import NPNet64
+import functools
+import random
+from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
 import torch
-
+import torch.nn as nn
+from einops import rearrange
+from torchvision.utils import make_grid
+import time
+from pytorch_lightning import seed_everything
+from torch import autocast
+from contextlib import contextmanager, nullcontext
+import accelerate
+import torchsde
+from SVDNoiseUnet import NPNet128
+from tqdm import tqdm, trange
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+model_repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"  # Replace to the model you would like to use
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+precision_scope = autocast
+   
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
 
+def append_zero(x):
+    return torch.cat([x, x.new_zeros([1])])
+
+# New helper to load a list-of-dicts preference JSON
+# JSON schema: [ { 'human_preference': [int], 'prompt': str, 'file_path': [str] }, ... ]
+def load_preference_json(json_path: str) -> list[dict]:
+    """Load records from a JSON file formatted as a list of preference dicts."""
+    with open(json_path, 'r') as f:
+        data = json.load(f)
+    return data
+
+# New helper to extract just the prompts from the preference JSON
+# Returns a flat list of all 'prompt' values
+
+def extract_prompts_from_pref_json(json_path: str) -> list[str]:
+    """Load a JSON of preference records and return only the prompts."""
+    records = load_preference_json(json_path)
+    return [rec['prompt'] for rec in records]
+
+# Example usage:
+# prompts = extract_prompts_from_pref_json("path/to/preference.json")
+# print(prompts)
+
+def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu',need_append_zero = True):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    ramp = torch.linspace(0, 1, n)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+    return append_zero(sigmas).to(device) if need_append_zero else sigmas.to(device)
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+def append_zero(x):
+    return torch.cat([x, x.new_zeros([1])])
+
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+
+class CFGDenoiser(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.inner_model = model
+
+    
+    def prepare_sdxl_pipeline_step_parameter(self, pipe, prompts, need_cfg, device):
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = pipe.encode_prompt(
+            prompt=prompts,
+            device=device,
+            do_classifier_free_guidance=need_cfg,
+        )
+    # timesteps = pipe.scheduler.timesteps
+    
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = pooled_prompt_embeds.to(device)
+        original_size = (1024, 1024)
+        crops_coords_top_left = (0, 0)
+        target_size = (1024, 1024)
+        text_encoder_projection_dim = None
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+        if pipe.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = pipe.text_encoder_2.config.projection_dim
+        passed_add_embed_dim = (
+            pipe.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+        )
+        expected_add_embed_dim = pipe.unet.add_embedding.linear_1.in_features
+        if expected_add_embed_dim != passed_add_embed_dim:
+            raise ValueError(
+                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+            )
+        add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype)
+        add_time_ids = add_time_ids.to(device)
+        negative_add_time_ids = add_time_ids
+
+        if need_cfg:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+        ret_dict = {
+            "text_embeds": add_text_embeds,
+            "time_ids": add_time_ids
+        }
+        return prompt_embeds, ret_dict
+
+    
+    def get_golden_noised(self, x, sigma,sigma_nxt, prompt, cond_scale,need_distill_uncond=False,tmp_list = [], uncond_list = [],noise_training_list={}):
+        x_in = torch.cat([x] * 2)
+        sigma_in = torch.cat([sigma] * 2)
+        sigma_nxt = torch.cat([sigma_nxt] * 2)
+        prompt_embeds, cond_kwargs = self.prepare_sdxl_pipeline_step_parameter(self.inner_model.pipe, prompt, need_cfg=True, device=self.inner_model.pipe.device)
+        _, ret = self.inner_model.get_customed_golden_noise(x_in
+                                                         , cond_scale
+                                                         , sigma_in, sigma_nxt
+                                                         , True
+                                                         , noise_training_list=noise_training_list
+                                                         , encoder_hidden_states=prompt_embeds.to(device=x.device, dtype=x.dtype)
+                                                         , added_cond_kwargs=cond_kwargs).chunk(2)
+
+        return ret
+
+    def forward(self, x, sigma, prompt, cond_scale,need_distill_uncond=False,tmp_list = [], uncond_list = []):
+        prompt_embeds, cond_kwargs = self.prepare_sdxl_pipeline_step_parameter(self.inner_model.pipe, prompt, need_cfg=True, device=self.inner_model.pipe.device)
+        # w = cond_scale * x.new_ones([x.shape[0]])
+        # w_embedding = guidance_scale_embedding(w, embedding_dim=self.inner_model.inner_model.config["time_cond_proj_dim"])
+        # w_embedding = w_embedding.to(device=x.device, dtype=x.dtype)
+        # # t = self.inner_model.sigma_to_t(sigma)
+        # cond = self.inner_model(
+        #     x,
+        #     sigma,
+        #     timestep_cond=w_embedding,
+        #     encoder_hidden_states=cond.to(device=x.device, dtype=x.dtype),
+        # )
+        # return cond 
+        x_in = torch.cat([x] * 2)
+        sigma_in = torch.cat([sigma] * 2)
+        # cond_in = torch.cat([uncond, cond])
+        uncond, cond = self.inner_model(x_in
+                                        , sigma_in
+                                        , tmp_list
+                                        , encoder_hidden_states=prompt_embeds.to(device=x.device, dtype=x.dtype)
+                                        , added_cond_kwargs=cond_kwargs).chunk(2)
+        if need_distill_uncond:
+            uncond_list.append(uncond)
+        return prompt_embeds, uncond + (cond - uncond) * cond_scale
+            
+    
+class DiscreteSchedule(nn.Module):
+    """A mapping between continuous noise levels (sigmas) and a list of discrete noise
+    levels."""
+
+    def __init__(self, sigmas, quantize):
+        super().__init__()
+        self.register_buffer('sigmas', sigmas)
+        self.register_buffer('log_sigmas', sigmas.log())
+        self.quantize = quantize
+
+    @property
+    def sigma_min(self):
+        return self.sigmas[0]
+
+    @property
+    def sigma_max(self):
+        return self.sigmas[-1]
+
+    def get_sigmas(self, n=None):
+        if n is None:
+            return append_zero(self.sigmas.flip(0))
+        t_max = len(self.sigmas) - 1
+        t = torch.linspace(t_max, 0, n, device=self.sigmas.device)
+        return append_zero(self.t_to_sigma(t))
+
+    def sigma_to_t(self, sigma, quantize=None):
+        quantize = self.quantize if quantize is None else quantize
+        log_sigma = sigma.log()
+        dists = log_sigma - self.log_sigmas[:, None]
+        if quantize:
+            return dists.abs().argmin(dim=0).view(sigma.shape)
+        low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+        low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx]
+        w = (low - log_sigma) / (low - high)
+        w = w.clamp(0, 1)
+        t = (1 - w) * low_idx + w * high_idx
+        return t.view(sigma.shape)
+
+    def t_to_sigma(self, t):
+        t = t.float()
+        low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
+        log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
+        return log_sigma.exp()
+
+class DiscreteEpsDDPMDenoiser(DiscreteSchedule):
+    """A wrapper for discrete schedule DDPM models that output eps (the predicted
+    noise)."""
+
+    def __init__(self, pipe, alphas_cumprod, quantize = False):
+        super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize)
+        self.pipe = pipe
+        self.inner_model = pipe.unet
+        # self.alphas_cumprod = alphas_cumprod.flip(0)
+        # Prepare a reversed version of alphas_cumprod for backward scheduling
+        self.register_buffer('alphas_cumprod', alphas_cumprod)
+        # self.register_buffer('alphas_cumprod_prev', append_zero(alphas_cumprod[:-1]))
+        self.sigma_data = 1.
+
+    def get_scalings(self, sigma):
+        c_out = -sigma
+        c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
+        return c_out, c_in
+
+    def get_eps(self, *args, **kwargs):
+        return self.inner_model(*args, **kwargs)
+
+    def get_alphact_and_sigma(self, timesteps, x_0, noise):
+        high_idx = torch.ceil(timesteps).int()
+        low_idx = torch.floor(timesteps).int()
+        
+        nxt_ts = timesteps - timesteps.new_ones(timesteps.shape[0])
+        
+        w = (timesteps - low_idx) / (high_idx - low_idx)
+        
+        beta_1 = torch.tensor([1e-4],dtype=torch.float32) 
+        beta_T = torch.tensor([0.02],dtype=torch.float32)
+        ddpm_max_step = torch.tensor([1000.0],dtype=torch.float32)
+        
+        beta_t: torch.Tensor = (beta_T - beta_1) / ddpm_max_step  * timesteps + beta_1
+        beta_t_prev: torch.Tensor = (beta_T - beta_1) / ddpm_max_step  * nxt_ts + beta_1
+        
+        alpha_t = beta_t.new_ones(beta_t.shape[0]) - beta_t
+        alpha_t_prev = beta_t.new_ones(beta_t.shape[0]) - beta_t_prev
+        
+        dir_xt = (1. - alpha_t_prev).sqrt() * noise
+        x_prev = alpha_t_prev.sqrt() * x_0 + dir_xt + noise
+        
+        alpha_cumprod_t_floor = self.alpha_cumprods[low_idx]
+        alpha_cumprod_t = (alpha_cumprod_t_floor * alpha_t) #.unsqueeze(1)
+        sqrt_alpha_cumprod_t = torch.sqrt(alpha_cumprod_t)
+        sigmas = torch.sqrt(alpha_cumprod_t.new_ones(alpha_cumprod_t.shape[0]) - alpha_cumprod_t)
+        
+        # Fix broadcasting
+        sqrt_alpha_cumprod_t = sqrt_alpha_cumprod_t[:, None, None]
+        sigmas = sigmas[:, None, None]
+        return alpha_cumprod_t, sigmas
+
+    def get_c_ins(self,sigmas): # use to adjust loss
+        ret = list()
+        for sigma in sigmas:
+            _, c_in = self.get_scalings(sigma=sigma)
+            ret.append(c_in)
+        return ret
+    
+    # def predicted_origin(model_output, timesteps, sample, alphas, sigmas, prediction_type = "epsilon"):
+    #     if prediction_type == "epsilon":
+    #         sigmas = extract_into_tensor(sigmas, timesteps, sample.shape)
+    #         alphas = extract_into_tensor(alphas, timesteps, sample.shape)
+    #         pred_x_0 = (sample - sigmas * model_output) / alphas
+    #     elif prediction_type == "v_prediction":
+    #         sigmas = extract_into_tensor(sigmas, timesteps, sample.shape)
+    #         alphas = extract_into_tensor(alphas, timesteps, sample.shape)
+    #         pred_x_0 = alphas * sample - sigmas * model_output
+    #     else:
+    #         raise ValueError(f"Prediction type {prediction_type} currently not supported.")
+    #     return pred_x_0
+    
+    def get_customed_golden_noise(self
+                                  , input
+                                  , unconditional_guidance_scale:float
+                                  , sigma
+                                  , sigma_nxt
+                                  ,  need_cond = True
+                                  , noise_training_list = {}
+                                  , **kwargs):
+        """User should ensure the input is a pure noise.  
+        It's a customed golden noise, not the one purposed in the paper.
+        Maybe the one purposed in the paper should be implemented in the future."""
+        c_out, c_in = [append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
+        
+        sigma_fn = lambda t: t.neg().exp()
+        t_fn = lambda sigma: sigma.log().neg()
+        if need_cond:
+            _, tmp_img = (input * c_in).chunk(2)
+        else :
+            tmp_img = input * c_in
+        # print(tmp_img.max())
+        # tmp_list.append(tmp_img)
+        eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs).sample
+        x_0 = input + eps * c_out
+        # normal_form_input = input * c_in
+        x_0_uncond, x_0 = x_0.chunk(2)
+        x_0 = x_0_uncond + unconditional_guidance_scale * (x_0 - x_0_uncond)
+        x_0 = torch.cat([x_0] * 2)
+        
+        
+        t, t_next = t_fn(sigma), t_fn(sigma_nxt)
+        h = t_next - t
+        
+        x = (append_dims(sigma_fn(t_next) / sigma_fn(t),input.ndim)) * input - append_dims((-h).expm1(),input.ndim) * x_0
+        
+        c_out_2, c_in_2 = [append_dims(x, input.ndim) for x in self.get_scalings(sigma_nxt)]
+        
+        
+        # e_t_uncond_ret, e_t_ret = self.get_eps(x * c_in_2, self.sigma_to_t(sigma_nxt), **kwargs).sample.chunk(2)
+        eps_ret = self.get_eps(x * c_in_2, self.sigma_to_t(sigma_nxt), **kwargs).sample
+        org_golden_noise = False
+        x_1 = x + eps_ret * c_out_2
+        if org_golden_noise:
+            ret = (x + append_dims((-h).expm1(),input.ndim) * x_1) / (append_dims(sigma_fn(t_next) / sigma_fn(t),input.ndim))
+        else :
+            e_t_uncond_ret, e_t_ret = eps_ret.chunk(2)
+            e_t_ret = e_t_uncond_ret + 1.0 * (e_t_ret - e_t_uncond_ret)
+            e_t_ret = torch.cat([e_t_ret] * 2)
+            ret = x_0 + e_t_ret * append_dims(sigma,input.ndim)
+        
+        noise_training_list['org_noise'] = input * c_in
+        noise_training_list['golden_noise'] = ret * c_in
+        # noise_training_list.append(tmp_dict)
+        return ret
+
+        # timesteps = self.sigma_to_t(sigma)
+
+        # high_idx = torch.ceil(timesteps).int().to(input.device)
+        # low_idx = torch.floor(timesteps).int().to(input.device)
+
+        # nxt_ts = (timesteps - timesteps.new_ones(timesteps.shape[0])).to(input.device) 
+        
+        # w = (timesteps - low_idx) / (high_idx - low_idx)
+        
+        # beta_1 = torch.tensor([1e-4],dtype=torch.float32).to(input.device)  
+        # beta_T = torch.tensor([0.02],dtype=torch.float32).to(input.device) 
+        # ddpm_max_step = torch.tensor([1000.0],dtype=torch.float32).to(input.device) 
+        
+        # beta_t: torch.Tensor = (beta_T - beta_1) / ddpm_max_step  * timesteps + beta_1
+        # beta_t_prev: torch.Tensor = (beta_T - beta_1) / ddpm_max_step  * nxt_ts + beta_1
+        
+        # alpha_t = beta_t.new_ones(beta_t.shape[0]) - beta_t
+        # alpha_t = append_dims(alpha_t, e_t.ndim)
+        # alpha_t_prev = beta_t_prev.new_ones(beta_t_prev.shape[0]) - beta_t_prev
+        # alpha_t_prev = append_dims(alpha_t_prev, e_t.ndim)
+        # alpha_cumprod_t_floor = self.alphas_cumprod[low_idx]
+        # alpha_cumprod_t_floor = append_dims(alpha_cumprod_t_floor, e_t.ndim)
+        # alpha_cumprod_t:torch.Tensor = (alpha_cumprod_t_floor * alpha_t) #.unsqueeze(1)
+        # alpha_cumprod_t_prev:torch.Tensor = (alpha_cumprod_t_floor * alpha_t_prev) #.unsqueeze(1)
+
+        # sqrt_one_minus_alphas = (1 - alpha_cumprod_t).sqrt()
+        
+        # dir_xt = (1. - alpha_cumprod_t_prev).sqrt() * e_t
+        # x_prev = alpha_cumprod_t_prev.sqrt() * x_0 + dir_xt
+        
+        # e_t_uncond_ret, e_t_ret = self.get_eps(x_prev, nxt_ts, **kwargs).sample.chunk(2)
+        # e_t_ret = e_t_uncond_ret + 1.0 * (e_t_ret - e_t_uncond_ret)
+        # e_t_ret = torch.cat([e_t_ret] * 2)
+        # x_ret = alpha_t.sqrt() * x_0 + sqrt_one_minus_alphas * e_t_ret
+        
+        # return x_ret
+    
+    def forward(self, input, sigma, tmp_list=[], need_cond = True, **kwargs):
+        # c_out_1, c_in_1 = [append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
+        # if need_cond:
+        #     tmp_img = input * c_in_1
+        # else :
+        #     tmp_img = input * c_in_1
+        # tmp_list.append(tmp_img)
+        # timestep = self.sigma_to_t(sigma)
+        # eps = self.get_eps(sample = input * c_in_1, timestep = timestep, **kwargs).sample
+        # c_skip, c_out = self.scalings_for_boundary_conditions(timestep=self.sigma_to_t(sigma))
+        # # return (input + eps * c_out_1) * c_out + input * c_in_1 * c_skip
+        # return (input + eps * c_out_1)
+        c_out, c_in = [append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
+        if need_cond:
+            _, tmp_img = (input * c_in).chunk(2)
+        else :
+            tmp_img = input * c_in
+        # print(tmp_img.max())
+        eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs).sample
+        tmp_x0 = input + eps * c_out 
+        tmp_dict = {'tmp_z': tmp_img, 'tmp_x0': tmp_x0}
+        tmp_list.append(tmp_dict)
+        return tmp_x0 #input + eps * c_out
+    
+    def get_special_sigmas_with_timesteps(self,timesteps):
+        low_idx, high_idx, w = np.minimum(np.floor(timesteps),999), np.minimum(np.ceil(timesteps),999), torch.from_numpy( timesteps - np.floor(timesteps))
+        self.alphas_cumprod = self.alphas_cumprod.to('cpu')
+        alphas = (1 - w) * self.alphas_cumprod[low_idx] + w * self.alphas_cumprod[high_idx]
+        return ((1 - alphas) / alphas) ** 0.5
+
+def get_ancestral_step(sigma_from, sigma_to, eta=1.):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    if not eta:
+        return sigma_to, 0.
+    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+    return sigma_down, sigma_up
+
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+class BatchedBrownianTree:
+    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
+    def __init__(self, x, t0, t1, seed=None, **kwargs):
+        t0, t1, self.sign = self.sort(t0, t1)
+        w0 = kwargs.get('w0', torch.zeros_like(x))
+        if seed is None:
+            seed = torch.randint(0, 2 ** 63 - 1, []).item()
+        self.batched = True
+        try:
+            assert len(seed) == x.shape[0]
+            w0 = w0[0]
+        except TypeError:
+            seed = [seed]
+            self.batched = False
+        self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
+
+    @staticmethod
+    def sort(a, b):
+        return (a, b, 1) if a < b else (b, a, -1)
+
+    def __call__(self, t0, t1):
+        t0, t1, sign = self.sort(t0, t1)
+        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
+        return w if self.batched else w[0]
+
+class BrownianTreeNoiseSampler:
+    """A noise sampler backed by a torchsde.BrownianTree.
+
+    Args:
+        x (Tensor): The tensor whose shape, device and dtype to use to generate
+            random samples.
+        sigma_min (float): The low end of the valid interval.
+        sigma_max (float): The high end of the valid interval.
+        seed (int or List[int]): The random seed. If a list of seeds is
+            supplied instead of a single integer, then the noise sampler will
+            use one BrownianTree per batch item, each with its own seed.
+        transform (callable): A function that maps sigma to the sampler's
+            internal timestep.
+    """
+    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
+        self.transform = transform
+        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
+        self.tree = BatchedBrownianTree(x, t0, t1, seed)
+
+    def __call__(self, sigma, sigma_next):
+        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
+        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+@torch.no_grad()
+def sample_euler(model
+                 , x
+                 , sigmas
+                 , extra_args=None
+                 , callback=None
+                 , disable=None
+                 , s_churn=0.
+                 , s_tmin=0.
+                 , s_tmax=float('inf')
+                 , tmp_list=[]
+                 , uncond_list=[]
+                 , need_distill_uncond=False
+                 , start_free_step = 1
+                 , noise_training_list={}
+                 , s_noise=1.):
+    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    intermediates = {'x_inter': [x],'pred_x0': []}
+    register_free_upblock2d(model.inner_model.pipe, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
+    register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if i == start_free_step:
+            register_free_upblock2d(model.inner_model.pipe, b1=1.3, b2=1.4, s1=0.9, s2=0.2)
+            register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.3, b2=1.4, s1=0.9, s2=0.2)
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        prompt_embeds, denoised = model(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+        d = to_d(x, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        dt = sigmas[i + 1] - sigma_hat
+        # Euler method
+        x = x + d * dt
+        intermediates['pred_x0'].append(denoised)
+        intermediates['x_inter'].append(x)
+    return prompt_embeds, intermediates, x
+
+@torch.no_grad()
+def sample_heun(model
+                , x
+                , sigmas
+                , extra_args=None
+                , callback=None
+                , disable=None
+                , s_churn=0.
+                , s_tmin=0.
+                , s_tmax=float('inf')
+                , tmp_list=[]
+                , uncond_list=[]
+                , need_distill_uncond=False
+                , noise_training_list={}
+                , s_noise=1.):
+    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    intermediates = {'x_inter': [x],'pred_x0': []}
+    register_free_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+    register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        prompt_embeds, denoised = model(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+        d = to_d(x, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        dt = sigmas[i + 1] - sigma_hat
+        if sigmas[i + 1] == 0:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            _, denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
+            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+            intermediates['pred_x0'].append(denoised_2)
+            intermediates['x_inter'].append(x)
+    return prompt_embeds, intermediates, x
+
+@torch.no_grad()
+def sample_dpmpp_ode(model
+                     , x
+                     , sigmas
+                     , need_golden_noise = False
+                     , start_free_step = 1
+                     , extra_args=None, callback=None
+                     , disable=None,tmp_list=[]
+                     , need_distill_uncond=False
+                     , need_raw_noise=False
+                     , uncond_list=[]
+                     , noise_training_list={}):
+    """DPM-Solver++."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    old_denoised = None
+    if need_raw_noise:
+        x = model.get_golden_noised(x=x,sigma=sigmas[0] * s_in, sigma_nxt=(sigmas[0] - 0.28) * s_in, noise_training_list=noise_training_list,**extra_args)
+    register_free_upblock2d(model.inner_model.pipe, b1=1, b2=1, s1=1, s2=1)
+    register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1, b2=1, s1=1, s2=1)
+    intermediates = {'x_inter': [x],'pred_x0': []}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if i == start_free_step:
+            register_free_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+            register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+        # macs, params = profile(model, inputs=(x, sigmas[i] * s_in,*extra_args.values(),need_distill_uncond,tmp_list,uncond_list, ))
+        prompt_embeds, denoised = model(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+        h = t_next - t
+            
+        x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
+        intermediates['pred_x0'].append(denoised)
+        intermediates['x_inter'].append(x)
+            
+        # print(denoised_d.max())
+        
+        # intermediates['noise'].append(denoised_d)
+    return prompt_embeds, intermediates,x
+
+@torch.no_grad()
+def sample_dpmpp_sde(model
+                     , x
+                     , sigmas
+                     , need_golden_noise = False
+                     , extra_args=None
+                     , callback=None
+                     , tmp_list=[]
+                     , need_distill_uncond=False
+                     , uncond_list=[]
+                     , disable=None, eta=1.
+                     , s_noise=1.
+                     , noise_sampler=None
+                     , r=1 / 2):
+    """DPM-Solver++ (stochastic)."""
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    if need_golden_noise:
+        x = model.get_golden_noised(x=x,sigma=sigmas[0] * s_in, sigma_nxt=sigmas[1] * s_in,**extra_args)
+    
+    intermediates = {'x_inter': [x],'pred_x0': []}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        prompt_embeds, denoised = model(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            # Euler method
+            d = to_d(x, sigmas[i], denoised)
+            dt = sigmas[i + 1] - sigmas[i]
+            x = x + d * dt
+            intermediates['pred_x0'].append(denoised)
+            intermediates['x_inter'].append(x)
+        else:
+            # DPM-Solver++
+            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+            h = t_next - t
+            s = t + h * r
+            fac = 1 / (2 * r)
+
+            # Step 1
+            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
+            s_ = t_fn(sd)
+            x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised
+            x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
+            prompt_embeds, denoised_2 = model(x_2, sigma_fn(s) * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args) #(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+
+            # Step 2
+            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
+            t_next_ = t_fn(sd)
+            denoised_d = (1 - fac) * denoised + fac * denoised_2
+            x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d
+            intermediates['pred_x0'].append(x)
+            x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
+            intermediates['x_inter'].append(x)
+    return  prompt_embeds, intermediates,x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m(model
+                    , x
+                    , sigmas
+                    # , need_golden_noise = True
+                    , extra_args=None
+                    , callback=None
+                    , disable=None
+                    , tmp_list=[]
+                    , need_distill_uncond=False
+                    , start_free_step=9
+                    , uncond_list=[]
+                    , stop_t = None):
+    """DPM-Solver++(2M)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    old_denoised = None
+    # if need_golden_noise:
+    #     x = model.get_golden_noised(x=x,sigma=sigmas[0] * s_in, sigma_nxt=sigmas[1] * s_in,**extra_args)
+    intermediates = {'x_inter': [x],'pred_x0': []}
+    register_free_upblock2d(model.inner_model.pipe, b1=1, b2=1, s1=1, s2=1)
+    register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1, b2=1, s1=1, s2=1)
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if i == start_free_step and len(sigmas) > 6:
+            register_free_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+            register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+        # else:
+        #     register_free_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=1.0, s2=1.0)
+        #     register_free_crossattn_upblock2d(model.inner_model.pipe, b1=1.1, b2=1.1, s1=1.0, s2=1.0)
+        # macs, params = profile(model, inputs=(x, sigmas[i] * s_in,*extra_args.values(),need_distill_uncond,tmp_list,uncond_list, ))
+        prompt_embeds, denoised = model(x, sigmas[i] * s_in, tmp_list=tmp_list,need_distill_uncond=need_distill_uncond,uncond_list=uncond_list, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+        h = t_next - t
+        if old_denoised is None or sigmas[i + 1] == 0:
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
+            intermediates['pred_x0'].append(denoised)
+            intermediates['x_inter'].append(x)
+        else:
+            h_last = t - t_fn(sigmas[i - 1])
+            r = h_last / h
+            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
+            intermediates['x_inter'].append(x)
+            intermediates['pred_x0'].append(denoised)
+            # print(denoised_d.max())
+        old_denoised = denoised
+        if i is not None and i == stop_t:
+            return intermediates, x
+        # intermediates['noise'].append(denoised_d)
+    return prompt_embeds, intermediates,x
+
+# Adapted from pipelines.StableDiffusionPipeline.encode_prompt
+def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompts, is_train=True):
+    captions = []
+    for caption in prompt_batch:
+        if random.random() < proportion_empty_prompts:
+            captions.append("")
+        elif isinstance(caption, str):
+            captions.append(caption)
+        elif isinstance(caption, (list, np.ndarray)):
+            # take a random caption if there are multiple
+            captions.append(random.choice(caption) if is_train else caption[0])
+
+    with torch.no_grad():
+        text_inputs = tokenizer(
+            captions,
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device))[0]
+
+    return prompt_embeds
+
+def chunk(it, size):
+    it = iter(it)
+    return iter(lambda: tuple(islice(it, size)), ())
+
+
+torch_dtype = torch.float32
+device = "cuda" 
+    # pipe = StableDiffusionPipeline.from_single_file( "./counterfeit/Counterfeit-V3.0_fp32.safetensors")
+repo_id = "madebyollin/sdxl-vae-fp16-fix"  # e.g., "distilbert/distilgpt2"
+filename = "sdxl_vae.safetensors"  # e.g., "pytorch_model.bin"
+downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename,cache_dir=".")
+
+    # pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')
+vae = AutoencoderKL.from_single_file(downloaded_path, torch_dtype=torch_dtype)
+vae.to('cuda')
+    
+pipe = StableDiffusionXLPipeline.from_pretrained("Lykon/dreamshaper-xl-1-0",torch_dtype=torch_dtype,vae=vae)
+    # pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",torch_dtype=torch.float16,vae=vae)
+    
+pipe.to('cuda')
+npn_net = NPNet128('SDXL', './sdxl.pth')
+
+pipe = pipe.to(device)
+register_free_upblock2d(pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+register_free_crossattn_upblock2d(pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
+noise_scheduler = pipe.scheduler
+alpha_schedule = noise_scheduler.alphas_cumprod.to(device=device, dtype=torch_dtype)
+model_wrap = DiscreteEpsDDPMDenoiser(pipe.unet, alpha_schedule, quantize=False)
+accelerator = accelerate.Accelerator()
+
+def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps):
+    """Helper function to generate image with specific number of steps"""
+    prompts =  [prompt]
+    if num_inference_steps <= 10:
+        register_free_upblock2d(pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+        register_free_crossattn_upblock2d(pipe, b1=1.1, b2=1.1, s1=0.9, s2=0.2)
+    else:
+        register_free_upblock2d(pipe, b1=1, b2=1, s1=1, s2=1)
+        register_free_crossattn_upblock2d(pipe, b1=1, b2=1, s1=1, s2=1)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        
+    def compute_embeddings(prompt_batch, proportion_empty_prompts, text_encoder, tokenizer, is_train=True):
+        prompt_embeds = encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompts, is_train)
+        return {"prompt_embeds": prompt_embeds}
+    
+    compute_embeddings_fn = functools.partial(
+        compute_embeddings,
+        proportion_empty_prompts=0,
+        text_encoder=pipe.text_encoder,
+        tokenizer=pipe.tokenizer,
+    )
+    generator = torch.Generator().manual_seed(seed)
+
+    intermediate_photos = list()
+                    # prompts = prompts[0]
+                            
+                    # if isinstance(prompts, tuple) or isinstance(prompts, str):
+                    #     prompts = list(prompts)
+    if isinstance(prompts, str):
+        prompts = prompts #+ 'high quality, best quality, masterpiece, 4K, highres, extremely detailed, ultra-detailed'
+        prompts = (prompts,)
+    if isinstance(prompts, tuple) or isinstance(prompts, str):
+        prompts = list(prompts)
+                    
+    shape = [4, height // 8, width // 8]
+    start_free_step = num_inference_steps
+    fir_stage_sigmas_ct = None
+    sec_stage_sigmas_ct = None
+                    # sigmas = model_wrap.get_sigmas(opt.ddim_steps).to(device=device)
+    if num_inference_steps == 5:
+        sigma_min, sigma_max = model_wrap.sigmas[0].item(), model_wrap.sigmas[-1].item()
+        sigmas = get_sigmas_karras(8, sigma_min, sigma_max, rho=5.0, device=device)# 6.0 if 5 else  10  10.0
+                    
+        ct_start, ct_end = model_wrap.sigma_to_t(sigmas[0]), model_wrap.sigma_to_t(sigmas[6])
+                        # sigma_kct_start, sigma_kct_end = sigmas[0].item(), sigmas[5].item()
+        ct = get_sigmas_karras(6, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
+        sigmas_ct = model_wrap.get_special_sigmas_with_timesteps(ct).to(device=device)
+    elif num_inference_steps == 6:
+        sigma_min, sigma_max = model_wrap.sigmas[0].item(), model_wrap.sigmas[-1].item()
+        sigmas = get_sigmas_karras(8, sigma_min, sigma_max,rho=5.0, device=device)# 6.0 if 5 else  10.0
+                    
+        ct_start, ct_end = model_wrap.sigma_to_t(sigmas[0]), model_wrap.sigma_to_t(sigmas[6])
+                        # sigma_kct_start, sigma_kct_end = sigmas[0].item(), sigmas[5].item()
+        ct = get_sigmas_karras(7, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
+        sigmas_ct = model_wrap.get_special_sigmas_with_timesteps(ct).to(device=device)
+        start_free_step = 6
+        fir_stage_sigmas_ct = sigmas_ct[:-2]
+        sec_stage_sigmas_ct = sigmas_ct[-3:]
+    elif num_inference_steps == 8:
+        sigma_min, sigma_max = model_wrap.sigmas[0].item(), model_wrap.sigmas[-1].item()
+        
+        sigmas = get_sigmas_karras(12, sigma_min, sigma_max,rho=12.0, device=device)# 6.0 if 5 else  10.0
+        ct_start, ct_end = model_wrap.sigma_to_t(sigmas[0]), model_wrap.sigma_to_t(sigmas[10])
+                        
+        ct = get_sigmas_karras(num_inference_steps + 1, ct_end.item(), ct_start.item(),rho=1.2, device='cpu',need_append_zero=False).numpy()
+        sigmas_ct = model_wrap.get_special_sigmas_with_timesteps(ct).to(device=device)
+        start_free_step = 8
+    else:
+        image = pipe(prompt=prompts
+            ,num_inference_steps=num_inference_steps
+            ,guidance_scale=guidance_scale
+            ,height=height
+            ,width=width).images[0]
+        return image
+    ts = []
+    for sigma in sigmas_ct:
+        t = model_wrap.sigma_to_t(sigma)
+        ts.append(t)
+                    
+    c_in = model_wrap.get_c_ins(sigmas=sigmas_ct)
+    x = torch.randn([1, *shape], device=device) * sigmas_ct[0]
+    model_wrap_cfg = CFGDenoiser(model_wrap)
+    (
+        c,
+        uc,
+        _,
+        _,
+    ) = pipe.encode_prompt(
+        prompt=prompts,
+        device=device,
+        do_classifier_free_guidance=True,
+    )
+                    # prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+
+    # if (num_inference_steps != -1 or num_inference_steps <= 8) and not opt.force_not_use_NPNet:
+    x = npn_net(x,c)
+                    
+    extra_args = {'prompt': prompts, 'cond_scale': guidance_scale}
+    with torch.no_grad():
+        # with precision_scope("cuda" if torch.cuda.is_available() else "cpu"):
+        if not (num_inference_steps == 8 or num_inference_steps == 7):
+            prompt_embeds, guide_distill, samples_ddim = sample_dpmpp_ode(model_wrap_cfg
+                                                                    , x
+                                                                    , fir_stage_sigmas_ct
+                                                                    , extra_args=extra_args
+                                                                    , disable=not accelerator.is_main_process
+                                                                    , need_raw_noise = False
+                                                                    , tmp_list=intermediate_photos)
+            _, _, samples_ddim = sample_euler(model_wrap_cfg
+                                                        , samples_ddim
+                                                        , sec_stage_sigmas_ct
+                                                        , extra_args=extra_args
+                                                        , disable=not accelerator.is_main_process
+                                                        , s_noise = 0.3
+                                                        , tmp_list=intermediate_photos)
+        else:
+            prompt_embeds, guide_distill, samples_ddim = sample_dpmpp_2m(model_wrap_cfg
+                                                                            , x
+                                                                            , sigmas_ct 
+                                                                            , extra_args=extra_args
+                                                                            , start_free_step=start_free_step
+                                                                            , disable=not accelerator.is_main_process
+                                                                            , tmp_list=intermediate_photos)
+                        # print('2m')
+                        
+        x_samples_ddim = pipe.vae.decode(samples_ddim / pipe.vae.config.scaling_factor).sample
+        x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
 
+        if True: # not opt.skip_save:
+            for x_sample in x_samples_ddim:
+                x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
+                image = Image.fromarray(x_sample.astype(np.uint8))
+            # base_count += 1
+    
+    # image = pipe(
+    #     prompt=prompt,
+    #     negative_prompt=negative_prompt,
+    #     guidance_scale=guidance_scale,
+    #     num_inference_steps=num_inference_steps,
+    #     width=width,
+    #     height=height,
+    #     generator=generator,
+    # ).images[0]
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
+    return image
+
+@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,
     negative_prompt,
     seed,
     randomize_seed,
-    width,
-    height,
+    resolution,
     guidance_scale,
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    
+    # Parse resolution string into width and height
+    width, height = map(int, resolution.split('x'))
+    
+    # Generate image with selected steps
+    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps)
+    
+    # Generate image with 50 steps for high quality
+    image_50_steps = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, 50)
 
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
+    return image_quick, image_50_steps, seed
 
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    "a painting of a virus monster playing guitar",
+    "a painting of a squirrel eating a burger",
 ]
 
 css = """
@@ -66,7 +972,7 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+        gr.Markdown(" # Hyperparameters are all you need")
 
         with gr.Row():
             prompt = gr.Text(
@@ -79,7 +985,13 @@ with gr.Blocks(css=css) as demo:
 
             run_button = gr.Button("Run", scale=0, variant="primary")
 
-        result = gr.Image(label="Result", show_label=False)
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Our fast inference Result")
+                result = gr.Image(label="Quick Result", show_label=False)
+            with gr.Column():
+                gr.Markdown("### Original 50 steps Result")
+                result_50_steps = gr.Image(label="50 Steps Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
@@ -99,22 +1011,15 @@ with gr.Blocks(css=css) as demo:
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
+            resolution = gr.Dropdown(
+                choices=[
+                    "1024x1024",
+                    "1216x832",
+                    "832x1216"
+                ],
+                value="1024x1024",
+                label="Resolution",
+            )
 
             with gr.Row():
                 guidance_scale = gr.Slider(
@@ -122,15 +1027,13 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=7.5,  # Replace with defaults that work for your model
                 )
 
-                num_inference_steps = gr.Slider(
+                num_inference_steps = gr.Dropdown(
+                    choices=[6, 8],
+                    value=8,
                     label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
@@ -142,12 +1045,11 @@ with gr.Blocks(css=css) as demo:
             negative_prompt,
             seed,
             randomize_seed,
-            width,
-            height,
+            resolution,
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed],
+        outputs=[result, result_50_steps, seed],
     )
 
 if __name__ == "__main__":

free_lunch_utils.py

@@ -0,0 +1,304 @@
+import torch
+import torch.fft as fft
+from diffusers.utils import is_torch_version
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+def isinstance_str(x: object, cls_name: str):
+    """
+    Checks whether x has any class *named* cls_name in its ancestry.
+    Doesn't require access to the class's implementation.
+    
+    Useful for patching!
+    """
+
+    for _cls in x.__class__.__mro__:
+        if _cls.__name__ == cls_name:
+            return True
+    
+    return False
+
+
+def Fourier_filter(x, threshold, scale):
+    dtype = x.dtype
+    x = x.type(torch.float32)
+    # FFT
+    x_freq = fft.fftn(x, dim=(-2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-2, -1))
+    
+    B, C, H, W = x_freq.shape
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    mask = torch.ones((B, C, H, W)).to(device=device)
+
+    crow, ccol = H // 2, W //2
+    mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale
+    x_freq = x_freq * mask
+
+    # IFFT
+    x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
+    x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
+    
+    x_filtered = x_filtered.type(dtype)
+    return x_filtered
+
+
+def register_upblock2d(model):
+    def up_forward(self):
+        def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
+            for resnet in self.resnets:
+                # pop res hidden states
+                res_hidden_states = res_hidden_states_tuple[-1]
+                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+                #print(f"in upblock2d, hidden states shape: {hidden_states.shape}")
+                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+                if self.training and self.gradient_checkpointing:
+
+                    def create_custom_forward(module):
+                        def custom_forward(*inputs):
+                            return module(*inputs)
+
+                        return custom_forward
+
+                    if is_torch_version(">=", "1.11.0"):
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
+                        )
+                    else:
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(resnet), hidden_states, temb
+                        )
+                else:
+                    hidden_states = resnet(hidden_states, temb)
+
+            if self.upsamplers is not None:
+                for upsampler in self.upsamplers:
+                    hidden_states = upsampler(hidden_states, upsample_size)
+
+            return hidden_states
+        
+        return forward
+    
+    for i, upsample_block in enumerate(model.unet.up_blocks):
+        if isinstance_str(upsample_block, "UpBlock2D"):
+            upsample_block.forward = up_forward(upsample_block)
+
+
+def register_free_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
+    def up_forward(self):
+        def forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
+            for resnet in self.resnets:
+                # pop res hidden states
+                res_hidden_states = res_hidden_states_tuple[-1]
+                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+                #print(f"in free upblock2d, hidden states shape: {hidden_states.shape}")
+                
+                # --------------- FreeU code -----------------------
+                # Only operate on the first two stages
+                if hidden_states.shape[1] == 1280:
+                    hidden_states[:,:640] = hidden_states[:,:640] * self.b1
+                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
+                if hidden_states.shape[1] == 640:
+                    hidden_states[:,:320] = hidden_states[:,:320] * self.b2
+                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
+                # ---------------------------------------------------------
+
+                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+                if self.training and self.gradient_checkpointing:
+
+                    def create_custom_forward(module):
+                        def custom_forward(*inputs):
+                            return module(*inputs)
+
+                        return custom_forward
+
+                    if is_torch_version(">=", "1.11.0"):
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
+                        )
+                    else:
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(resnet), hidden_states, temb
+                        )
+                else:
+                    hidden_states = resnet(hidden_states, temb)
+
+            if self.upsamplers is not None:
+                for upsampler in self.upsamplers:
+                    hidden_states = upsampler(hidden_states, upsample_size)
+
+            return hidden_states
+        
+        return forward
+    
+    for i, upsample_block in enumerate(model.unet.up_blocks):
+        if isinstance_str(upsample_block, "UpBlock2D"):
+            upsample_block.forward = up_forward(upsample_block)
+            setattr(upsample_block, 'b1', b1)
+            setattr(upsample_block, 'b2', b2)
+            setattr(upsample_block, 's1', s1)
+            setattr(upsample_block, 's2', s2)
+
+
+def register_crossattn_upblock2d(model):
+    def up_forward(self):
+        def forward(
+            hidden_states: torch.FloatTensor,
+            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
+            temb: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.FloatTensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            upsample_size: Optional[int] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        ):
+            for resnet, attn in zip(self.resnets, self.attentions):
+                # pop res hidden states
+                #print(f"in crossatten upblock2d, hidden states shape: {hidden_states.shape}")
+                res_hidden_states = res_hidden_states_tuple[-1]
+                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+                if self.training and self.gradient_checkpointing:
+
+                    def create_custom_forward(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(resnet),
+                        hidden_states,
+                        temb,
+                        **ckpt_kwargs,
+                    )
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(attn, return_dict=False),
+                        hidden_states,
+                        encoder_hidden_states,
+                        None,  # timestep
+                        None,  # class_labels
+                        cross_attention_kwargs,
+                        attention_mask,
+                        encoder_attention_mask,
+                        **ckpt_kwargs,
+                    )[0]
+                else:
+                    hidden_states = resnet(hidden_states, temb)
+                    hidden_states = attn(
+                        hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        attention_mask=attention_mask,
+                        encoder_attention_mask=encoder_attention_mask,
+                        return_dict=False,
+                    )[0]
+
+            if self.upsamplers is not None:
+                for upsampler in self.upsamplers:
+                    hidden_states = upsampler(hidden_states, upsample_size)
+
+            return hidden_states
+        
+        return forward
+    
+    for i, upsample_block in enumerate(model.unet.up_blocks):
+        if isinstance_str(upsample_block, "CrossAttnUpBlock2D"):
+            upsample_block.forward = up_forward(upsample_block)
+
+
+def register_free_crossattn_upblock2d(model, b1=1.2, b2=1.4, s1=0.9, s2=0.2):
+    def up_forward(self):
+        def forward(
+            hidden_states: torch.FloatTensor,
+            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
+            temb: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.FloatTensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            upsample_size: Optional[int] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        ):
+            for resnet, attn in zip(self.resnets, self.attentions):
+                # pop res hidden states
+                #print(f"in free crossatten upblock2d, hidden states shape: {hidden_states.shape}")
+                res_hidden_states = res_hidden_states_tuple[-1]
+                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+
+                # --------------- FreeU code -----------------------
+                # Only operate on the first two stages
+                if hidden_states.shape[1] == 1280:
+                    hidden_states[:,:640] = hidden_states[:,:640] * self.b1
+                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s1)
+                if hidden_states.shape[1] == 640:
+                    hidden_states[:,:320] = hidden_states[:,:320] * self.b2
+                    res_hidden_states = Fourier_filter(res_hidden_states, threshold=1, scale=self.s2)
+                # ---------------------------------------------------------
+
+                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+                if self.training and self.gradient_checkpointing:
+
+                    def create_custom_forward(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(resnet),
+                        hidden_states,
+                        temb,
+                        **ckpt_kwargs,
+                    )
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(attn, return_dict=False),
+                        hidden_states,
+                        encoder_hidden_states,
+                        None,  # timestep
+                        None,  # class_labels
+                        cross_attention_kwargs,
+                        attention_mask,
+                        encoder_attention_mask,
+                        **ckpt_kwargs,
+                    )[0]
+                else:
+                    hidden_states = resnet(hidden_states, temb)
+                    # hidden_states = attn(
+                    #     hidden_states,
+                    #     encoder_hidden_states=encoder_hidden_states,
+                    #     cross_attention_kwargs=cross_attention_kwargs,
+                    #     encoder_attention_mask=encoder_attention_mask,
+                    #     return_dict=False,
+                    # )[0]
+                    hidden_states = attn(
+                        hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                    )[0]
+
+            if self.upsamplers is not None:
+                for upsampler in self.upsamplers:
+                    hidden_states = upsampler(hidden_states, upsample_size)
+
+            return hidden_states
+        
+        return forward
+    
+    for i, upsample_block in enumerate(model.unet.up_blocks):
+        if isinstance_str(upsample_block, "CrossAttnUpBlock2D"):
+            upsample_block.forward = up_forward(upsample_block)
+            setattr(upsample_block, 'b1', b1)
+            setattr(upsample_block, 'b2', b2)
+            setattr(upsample_block, 's1', s1)
+            setattr(upsample_block, 's2', s2)

requirements.txt

@@ -1,6 +1,10 @@
-accelerate
-diffusers
-invisible_watermark
-torch
-transformers
-xformers
+tqdm
+einops
+pytorch_lightning
+accelerate>=0.20.0
+torchsde
+pycocotools
+diffusers== 0.32.2
+timm
+transformers==4.49
+torch>=2.0.0

sdxl.pth

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