acmyu/KeyframesAI

1import glob
2import os
3import torch
4from torch import nn
5from PIL import Image
6import numpy as np
7from diffusers import UniPCMultistepScheduler
8from src.models.stage2_inpaint_unet_2d_condition import Stage2_InapintUNet2DConditionModel
9
10from torchvision import transforms
11from diffusers.models.controlnet import ControlNetConditioningEmbedding
12from transformers import CLIPImageProcessor
13from transformers import Dinov2Model
14from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel,ControlNetModel,DDIMScheduler
15from src.pipelines.PCDMs_pipeline import PCDMsPipeline
16#from single_extract_pose import inference_pose
17
18
19class ImageProjModel_p(torch.nn.Module):
20    """SD model with image prompt"""
21
22    def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):
23        super().__init__()
24
25        self.net = nn.Sequential(
26            nn.Linear(in_dim, hidden_dim),
27            nn.GELU(),
28            nn.Dropout(dropout),
29            nn.LayerNorm(hidden_dim),
30            nn.Linear(hidden_dim, out_dim),
31            nn.Dropout(dropout)
32        )
33
34    def forward(self, x): 
35        return self.net(x)
36
37class ImageProjModel(torch.nn.Module):
38    """SD model with image prompt"""
39    def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):
40        super().__init__()
41
42        self.net = nn.Sequential(
43            nn.Linear(in_dim, hidden_dim),
44            nn.GELU(),
45            nn.Dropout(dropout),
46            nn.LayerNorm(hidden_dim),
47            nn.Linear(hidden_dim, out_dim),
48            nn.Dropout(dropout)
49        )
50
51    def forward(self, x):  
52        return self.net(x)
53        
54class SDModel(torch.nn.Module):
55    """SD model with image prompt"""
56    def __init__(self, unet) -> None:
57        super().__init__()
58        self.image_proj_model = ImageProjModel(in_dim=1536, hidden_dim=768, out_dim=1024).to(device).to(dtype=torch.float16)
59        self.unet = unet
60        self.pose_proj = ControlNetConditioningEmbedding(
61            conditioning_embedding_channels=320,
62            block_out_channels=(16, 32, 96, 256),
63            conditioning_channels=3).to(device).to(dtype=torch.float16)
64
65
66def image_grid(imgs, rows, cols):
67    assert len(imgs) == rows * cols
68    w, h = imgs[0].size
69    print(w, h)
70    grid = Image.new("RGB", size=(cols * w, rows * h))
71    grid_w, grid_h = grid.size
72
73    for i, img in enumerate(imgs):
74        grid.paste(img, box=(i % cols * w, i // cols * h))
75    return grid
76
77def load_mydict(model_ckpt_path):
78    model_sd = torch.load(model_ckpt_path, map_location="cpu")["module"]
79
80    image_proj_model_dict = {}
81    pose_proj_dict = {}
82    unet_dict = {}
83    for k in model_sd.keys():
84        if k.startswith("pose_proj"):
85            pose_proj_dict[k.replace("pose_proj.", "")] = model_sd[k]
86
87        elif k.startswith("image_proj_model_p"):
88            image_proj_model_dict[k.replace("image_proj_model_p.", "")] = model_sd[k]
89        elif k.startswith("image_proj_model"):
90            image_proj_model_dict[k.replace("image_proj_model.", "")] = model_sd[k]
91
92
93        elif k.startswith("unet"):
94            unet_dict[k.replace("unet.", "")] = model_sd[k]
95        else:
96            print(k)
97    return image_proj_model_dict, pose_proj_dict, unet_dict
98
99
100        
101device = "cuda"
102pretrained_model_name_or_path ="stabilityai/stable-diffusion-2-1-base"
103image_encoder_path = "facebook/dinov2-giant"
104model_ckpt_path = "./pcdms_ckpt.pt"   # ckpt path
105#model_ckpt_path = 'fine_tuned_pcdms.pt'
106
107
108clip_image_processor = CLIPImageProcessor()
109img_transform = transforms.Compose([
110    transforms.ToTensor(),
111    transforms.Normalize([0.5], [0.5]),
112])
113
114generator = torch.Generator(device=device).manual_seed(42)
115unet = Stage2_InapintUNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, torch_dtype=torch.float16,subfolder="unet",in_channels=9, low_cpu_mem_usage=False, ignore_mismatched_sizes=True).to(device)
116vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path,subfolder="vae").to(device, dtype=torch.float16)
117image_encoder = Dinov2Model.from_pretrained(image_encoder_path).to(device, dtype=torch.float16)
118noise_scheduler = DDIMScheduler(
119    num_train_timesteps=1000,
120    beta_start=0.00085,
121    beta_end=0.012,
122    beta_schedule="scaled_linear",
123    clip_sample=False,
124    set_alpha_to_one=False,
125    steps_offset=1,
126)
127
128image_proj_model = ImageProjModel(in_dim=1536, hidden_dim=768, out_dim=1024).to(device).to(dtype=torch.float16)
129pose_proj_model = ControlNetConditioningEmbedding(
130    conditioning_embedding_channels=320,
131    block_out_channels=(16, 32, 96, 256),
132    conditioning_channels=3).to(device).to(dtype=torch.float16)
133
134
135# load weight
136image_proj_model_dict, pose_proj_dict, unet_dict = load_mydict(model_ckpt_path)
137image_proj_model.load_state_dict(image_proj_model_dict)
138pose_proj_model.load_state_dict(pose_proj_dict)
139unet.load_state_dict(unet_dict)
140
141
142pipe = PCDMsPipeline.from_pretrained(pretrained_model_name_or_path, unet=unet,  torch_dtype=torch.float16, scheduler=noise_scheduler,feature_extractor=None,safety_checker=None).to(device)
143
144print('====================== model load finish ===================')
145
146
147
148
149for n in range(1, 13):
150
151    num_samples = 1
152    image_size = (512, 512)
153    s_img_path = 'imgs/sm.png' # input image 1
154    target_pose_img = 'imgs/walk'+str(n)+'.png' # input image 2
155    
156    #t_pose = inference_pose(target_pose_img, image_size=(image_size[1], image_size[0])).resize(image_size, Image.BICUBIC)
157    t_pose = Image.open(target_pose_img).convert("RGB").resize((image_size), Image.BICUBIC)
158
159
160    s_img = Image.open(s_img_path).convert("RGB").resize(image_size, Image.BICUBIC)
161    black_image = Image.new("RGB", s_img.size, (0, 0, 0)).resize(image_size, Image.BICUBIC)
162
163    s_img_t_mask = Image.new("RGB", (s_img.width * 2, s_img.height))
164    s_img_t_mask.paste(s_img, (0, 0))
165    s_img_t_mask.paste(black_image, (s_img.width, 0))
166
167    #s_pose = inference_pose(s_img_path, image_size=(image_size[1], image_size[0])).resize(image_size, Image.BICUBIC)
168    s_pose = Image.open('imgs/sm_pose.jpg').convert("RGB").resize(image_size, Image.BICUBIC)
169    print('source image width: {}, height: {}'.format(s_pose.width, s_pose.height))
170    #t_pose = Image.open(target_pose_img).convert("RGB").resize((image_size), Image.BICUBIC)
171
172    st_pose = Image.new("RGB", (s_pose.width * 2, s_pose.height))
173    st_pose.paste(s_pose, (0, 0))
174    st_pose.paste(t_pose, (s_pose.width, 0))
175
176
177    clip_s_img = clip_image_processor(images=s_img, return_tensors="pt").pixel_values
178    vae_image = torch.unsqueeze(img_transform(s_img_t_mask), 0)
179    cond_st_pose = torch.unsqueeze(img_transform(st_pose), 0)
180
181    mask1 = torch.ones((1, 1, int(image_size[0] / 8), int(image_size[1] / 8))).to(device, dtype=torch.float16)
182    mask0 = torch.zeros((1, 1, int(image_size[0] / 8), int(image_size[1] / 8))).to(device, dtype=torch.float16)
183    mask = torch.cat([mask1, mask0], dim=3)
184
185
186    with torch.inference_mode():
187        cond_pose = pose_proj_model(cond_st_pose.to(dtype=torch.float16, device=device))
188        simg_mask_latents = pipe.vae.encode(vae_image.to(device, dtype=torch.float16)).latent_dist.sample()
189        simg_mask_latents = simg_mask_latents * 0.18215
190
191        images_embeds = image_encoder(clip_s_img.to(device, dtype=torch.float16)).last_hidden_state
192        image_prompt_embeds = image_proj_model(images_embeds)
193        uncond_image_prompt_embeds = image_proj_model(torch.zeros_like(images_embeds))
194
195    bs_embed, seq_len, _ = image_prompt_embeds.shape
196    image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
197    image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
198    uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
199    uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
200
201    output, _ = pipe(
202        simg_mask_latents= simg_mask_latents,
203        mask = mask,
204        cond_pose = cond_pose,
205        prompt_embeds=image_prompt_embeds,
206        negative_prompt_embeds=uncond_image_prompt_embeds,
207        height=image_size[1],
208        width=image_size[0]*2,
209        num_images_per_prompt=num_samples,
210        guidance_scale=2.0,
211        generator=generator,
212        num_inference_steps=50,
213    )
214
215    output = output.images[-1]
216    
217    result = output.crop((image_size[0], 0, image_size[0] * 2, image_size[1]))
218    result.resize((352, 512), Image.BICUBIC).save(str(n)+'.png')
219
220    """
221    save_output = []
222    result = output.crop((image_size[0], 0, image_size[0] * 2, image_size[1]))
223    save_output.append(result.resize((352, 512), Image.BICUBIC))
224    save_output.insert(0, t_pose.resize((352, 512), Image.BICUBIC))
225    save_output.insert(0, s_img.resize((352, 512), Image.BICUBIC))
226    grid = image_grid(save_output, 1, 3)
227    grid.save("out.png")
228    """
229