add inference code

README.md

@@ -1,3 +1,91 @@
 ---
 license: apache-2.0
 ---
+
+# NPNet Pipeline Usage Guide😄 
+
+## Overview
+
+This guide provides instructions on how to use the NPNet, a noise prompt network aims to transform the random Gaussian noise into golden noise, by adding a small desirable perturbation derived from the text prompt to boost the overall quality and semantic faithfulness of the synthesized images.
+
+Here we provide the inference code which supports different models like ***Stable Diffusion XL, DreamShaper-xl-v2-turbo, and Hunyuan-DiT.***
+
+
+## Requirements
+
+- `python >= 3.8.0`
+- `pytorch with cuda version`
+- `diffusers`
+- `PIL`
+- `numpy`
+- `timm`
+- `argparse`
+- `einops`
+
+## Installation🚀️
+
+Make sure you have successfully built `python` environment and installed `pytorch` with cuda version. Before running the script, ensure you have all the required packages installed. You can install them using:
+
+```bash
+pip install diffusers, PIL, numpy, timm, argparse, einops
+```
+
+## Usage👀️ 
+
+To use the NPNet pipeline, you need to run the `npnet_pipeline.py` script with appropriate command-line arguments. Below are the available options:
+
+### Command-Line Arguments
+
+- `--pipeline`: Select the model pipeline (`SDXL`, `DreamShaper`, `DiT`). Default is `SDXL`.
+- `--prompt`: The textual prompt based on which the image will be generated. Default is "A banana on the left of an apple."
+- `--inference-step`: Number of inference steps for the diffusion process. Default is 50.
+- `--cfg`: Classifier-free guidance scale. Default is 5.5.
+- `--pretrained-path`: Path to the pretrained model weights. Default is a specified path in the script.
+- `--size`: The size (height and width) of the generated image. Default is 1024.
+
+### Running the Script
+
+Run the script from the command line by navigating to the directory containing `npnet_pipeline.py` and executing:
+
+```
+python npnet_pipeline.py --pipeline SDXL --prompt "A banana on the left of an apple." --size 1024
+```
+
+This command will generate an image based on the prompt "A banana on the left of an apple." using the Stable Diffusion XL model with an image size of 1024x1024 pixels.
+
+### Output🎉️ 
+
+The script will save two images:
+
+- A standard image generated by the diffusion model.
+- A golden image generated by the diffusion model with the NPNet.
+
+Both images will be saved in the current directory with names based on the model and prompt.
+
+## Pre-trained Weights Download❤️
+
+We provide the pre-trained NPNet weights of Stable Diffusion XL, DreamShaper-xl-v2-turbo, and Hunyuan-DiT with [google drive](https://drive.google.com/drive/folders/1Z0wg4HADhpgrztyT3eWijPbJJN5Y2jQt?usp=drive_link)
+
+## Citation:
+If you find our code useful for your research, please cite our paper.
+
+```
+@misc{zhou2024goldennoisediffusionmodels,
+      title={Golden Noise for Diffusion Models: A Learning Framework}, 
+      author={Zikai Zhou and Shitong Shao and Lichen Bai and Zhiqiang Xu and Bo Han and Zeke Xie},
+      year={2024},
+      eprint={2411.09502},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2411.09502}, 
+}
+```
+
+## 🙏 Acknowledgements
+
+We thank the community and contributors for their invaluable support in developing NPNet. 
+We thank @DataCTE for constructing the ComfyUI of NPNet inference code [ComfyUI](https://github.com/DataCTE/ComfyUI_Golden-Noise).
+We thank @asagi4 for constructing the ComfyUI of NPNet inference code [ComfyUI](https://github.com/asagi4/ComfyUI-NPNet).
+
+---
+

model/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):
+        super().__init__()
+        self.upsample = lambda x: F.interpolate(x, [224,224])
+        self.downsample = lambda x: F.interpolate(x, [resolution,resolution])
+        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

model/SVDNoiseUnet.py

@@ -0,0 +1,109 @@
+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
+
+__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): # resolution = size // 8
+        super(SVDNoiseUnet, 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=128):
+        super(SVDNoiseUnet_Concise, self).__init__()
+        

model/__init__.py

@@ -0,0 +1,2 @@
+from .NoiseTransformer import NoiseTransformer
+from .SVDNoiseUnet import SVDNoiseUnet

npnet_pipeline.py

@@ -0,0 +1,168 @@
+import random
+import torch
+import argparse
+import torch.nn as nn
+import numpy as np
+
+from PIL import Image
+from diffusers.models.normalization import AdaGroupNorm
+from diffusers import  DDIMScheduler, DPMSolverMultistepScheduler, \
+                  DDPMScheduler, StableDiffusionXLPipeline, HunyuanDiTPipeline
+                        
+                    
+from model import NoiseTransformer, SVDNoiseUnet
+
+
+class NPNet(nn.Module):
+      def __init__(self, model_id, pretrained_path=True, device='cuda') -> None:
+            super(NPNet, 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).to(self.device).to(torch.float32)
+            unet_svd = SVDNoiseUnet(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
+      
+
+def get_args():
+      parser = argparse.ArgumentParser()
+
+      # model and dataset construction
+      parser.add_argument('--pipeline', default='SDXL', 
+                        choices=['SDXL', 'DreamShaper', 'DiT'], type=str)
+      parser.add_argument('--prompt', default='A banana on the left of an apple.', type=str)
+      parser.add_argument("--inference-step", default=50, type=int)
+
+      # for dreamershaper is 3.5, remaining is 5.5, DiT is 5.0
+      parser.add_argument("--cfg", default=5.5, type=float)
+
+      # model pretrained weight path
+      parser.add_argument('--pretrained-path', type=str,
+                        default='xxx')
+
+      parser.add_argument("--size", default=1024, type=int)
+
+      args = parser.parse_args()
+
+      print("generating config:")
+      print(f"Config: {args}")
+      print('-' * 100)
+
+      return args
+
+
+def main(args):
+      dtype = torch.float16
+      device = torch.device('cuda')
+
+      if args.pipeline == 'SDXL':
+
+            pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
+                                                            variant="fp16",use_safetensors=True,
+                                                            torch_dtype=torch.float16).to(device)
+            
+      elif args.pipeline == 'DreamShaper':
+            pipe = StableDiffusionXLPipeline.from_pretrained("lykon/dreamshaper-xl-v2-turbo",
+                                                            torch_dtype=torch.float16, 
+                                                            variant="fp16").to(device)
+      
+      else:
+             pipe = HunyuanDiTPipeline.from_pretrained("Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", 
+                                                            torch_dtype=torch.float16).to(device)
+             
+      pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+      pipe.enable_model_cpu_offload()
+
+      # create the initial noise
+      latent = torch.randn(1, 4, 128, 128, dtype=dtype).to(device)
+
+
+      # use the pre-trained text encoder in T2I models to encode prompts
+      prompt_embeds, _, _, _= pipe.encode_prompt(prompt=args.prompt, device=device)
+
+      # create NPNet to get the target noise
+      npn_net = NPNet(args.pipeline, args.pretrained_path)
+
+      golden_noise = npn_net(latent, prompt_embeds)
+
+      # standard inference pipeline
+      latent = latent.half()
+      golden_noise = golden_noise.half()
+
+      pipe = pipe.to(torch.float16)
+
+      standard_img = pipe(
+            prompt=args.prompt,
+            height=args.size,
+            width=args.size,
+            num_inference_steps=args.inference_step,
+            guidance_scale=args.cfg,
+            latents=latent).images[0]
+      
+      golden_img = pipe(
+            prompt=args.prompt,
+            height=args.size,
+            width=args.size,
+            num_inference_steps=args.inference_step,
+            guidance_scale=args.cfg,
+            latents=golden_noise).images[0]
+      
+      # image save path
+      standard_img.save(f"{args.pipeline}_{args.prompt}_standard_image.jpg")
+      golden_img.save(f"{args.pipeline}_{args.prompt}_golden_image.jpg")
+      
+
+if __name__ == '__main__':
+      args = get_args()
+      main(args)
+      

weights/dit.pth

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

weights/dreamshaper.pth

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

weights/sdxl.pth

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