File size: 6,069 Bytes
949adbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 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)
|