import glob import os import sys import copy from typing import List import numpy as np import torch from einops import rearrange from omegaconf import OmegaConf from PIL import Image from pytorch_lightning import seed_everything from pytorch3d.renderer.cameras import PerspectiveCameras from pytorch3d.renderer import look_at_view_transform from pytorch3d.renderer.camera_utils import join_cameras_as_batch import json sys.path.append('./') from sgm.util import instantiate_from_config, load_safetensors choices = [] def load_base_model(config, ckpt=None, verbose=True): config = OmegaConf.load(config) # load model config.model.params.network_config.params.far = 3 config.model.params.first_stage_config.params.ckpt_path = "pretrained-models/sdxl_vae.safetensors" guider_config = {'target': 'sgm.modules.diffusionmodules.guiders.ScheduledCFGImgTextRef', 'params': {'scale': 7.5, 'scale_im': 3.5} } config.model.params.sampler_config.params.guider_config = guider_config model = instantiate_from_config(config.model) if ckpt is not None: print(f"Loading model from {ckpt}") if ckpt.endswith("ckpt"): pl_sd = torch.load(ckpt, map_location="cpu") if "global_step" in pl_sd: print(f"Global Step: {pl_sd['global_step']}") sd = pl_sd["state_dict"] elif ckpt.endswith("safetensors"): sd = load_safetensors(ckpt) if 'modifier_token' in config.data.params: del sd['conditioner.embedders.0.transformer.text_model.embeddings.token_embedding.weight'] del sd['conditioner.embedders.1.model.token_embedding.weight'] else: raise NotImplementedError m, u = model.load_state_dict(sd, strict=False) model.cuda() model.eval() return model def load_delta_model(model, delta_ckpt=None, verbose=True, freeze=True): """ model is preloaded base stable diffusion model """ msg = None if delta_ckpt is not None: pl_sd_delta = torch.load(delta_ckpt, map_location="cpu") sd_delta = pl_sd_delta["delta_state_dict"] # TODO: add new delta loading embedding stuff? for name, module in model.model.diffusion_model.named_modules(): if len(name.split('.')) > 1 and name.split('.')[-2] == 'transformer_blocks': if hasattr(module, 'pose_emb_layers'): module.register_buffer('references', sd_delta[f'model.diffusion_model.{name}.references']) del sd_delta[f'model.diffusion_model.{name}.references'] m, u = model.load_state_dict(sd_delta, strict=False) if len(m) > 0 and verbose: print("missing keys:") if len(u) > 0 and verbose: print("unexpected keys:") if freeze: for param in model.parameters(): param.requires_grad = False model.cuda() model.eval() return model, msg def get_unique_embedder_keys_from_conditioner(conditioner): p = [x.input_keys for x in conditioner.embedders] return list(set([item for sublist in p for item in sublist])) + ['jpg_ref'] def customforward(self, x, xr, context=None, contextr=None, pose=None, mask_ref=None, prev_weights=None, timesteps=None, drop_im=None): # note: if no context is given, cross-attention defaults to self-attention if not isinstance(context, list): context = [context] b, c, h, w = x.shape x_in = x fg_masks = [] alphas = [] rgbs = [] x = self.norm(x) if not self.use_linear: x = self.proj_in(x) x = rearrange(x, "b c h w -> b (h w) c").contiguous() if self.use_linear: x = self.proj_in(x) prev_weights = None counter = 0 for i, block in enumerate(self.transformer_blocks): if i > 0 and len(context) == 1: i = 0 # use same context for each block if self.image_cross and (counter % self.poscontrol_interval == 0): x, fg_mask, weights, alpha, rgb = block(x, context=context[i], context_ref=x, pose=pose, mask_ref=mask_ref, prev_weights=prev_weights, drop_im=drop_im) prev_weights = weights fg_masks.append(fg_mask) if alpha is not None: alphas.append(alpha) if rgb is not None: rgbs.append(rgb) else: x, _, _, _, _ = block(x, context=context[i], drop_im=drop_im) counter += 1 if self.use_linear: x = self.proj_out(x) x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous() if not self.use_linear: x = self.proj_out(x) if len(fg_masks) > 0: if len(rgbs) <= 0: rgbs = None if len(alphas) <= 0: alphas = None return x + x_in, None, fg_masks, prev_weights, alphas, rgbs else: return x + x_in, None, None, prev_weights, None, None def _customforward( self, x, context=None, context_ref=None, pose=None, mask_ref=None, prev_weights=None, drop_im=None, additional_tokens=None, n_times_crossframe_attn_in_self=0 ): if context_ref is not None: global choices batch_size = x.size(0) # IP2P like sampling or default sampling if batch_size % 3 == 0: batch_size = batch_size // 3 context_ref = torch.stack([self.references[:-1][y] for y in choices]).unsqueeze(0).expand(batch_size, -1, -1, -1) context_ref = torch.cat([self.references[-1:].unsqueeze(0).expand(batch_size, context_ref.size(1), -1, -1), context_ref, context_ref], dim=0) else: batch_size = batch_size // 2 context_ref = torch.stack([self.references[:-1][y] for y in choices]).unsqueeze(0).expand(batch_size, -1, -1, -1) context_ref = torch.cat([self.references[-1:].unsqueeze(0).expand(batch_size, context_ref.size(1), -1, -1), context_ref], dim=0) fg_mask = None weights = None alphas = None predicted_rgb = None x = ( self.attn1( self.norm1(x), context=context if self.disable_self_attn else None, additional_tokens=additional_tokens, n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self if not self.disable_self_attn else 0, ) + x ) x = ( self.attn2( self.norm2(x), context=context, additional_tokens=additional_tokens, ) + x ) if context_ref is not None: if self.rendered_feat is not None: x = self.pose_emb_layers(torch.cat([x, self.rendered_feat], dim=-1)) else: xref, fg_mask, weights, alphas, predicted_rgb = self.reference_attn(x, context_ref, context, pose, prev_weights, mask_ref) self.rendered_feat = xref x = self.pose_emb_layers(torch.cat([x, xref], -1)) x = self.ff(self.norm3(x)) + x return x, fg_mask, weights, alphas, predicted_rgb def log_images( model, batch, N: int = 1, noise=None, scale_im=3.5, num_steps: int = 10, ucg_keys: List[str] = None, **kwargs, ): log = dict() conditioner_input_keys = [e.input_keys for e in model.conditioner.embedders] ucg_keys = conditioner_input_keys pose = batch['pose'] c, uc = model.conditioner.get_unconditional_conditioning( batch, force_uc_zero_embeddings=ucg_keys if len(model.conditioner.embedders) > 0 else [], force_ref_zero_embeddings=True ) _, n = 1, len(pose)-1 sampling_kwargs = {} if scale_im > 0: if uc is not None: if isinstance(pose, list): pose = pose[:N]*3 else: pose = torch.cat([pose[:N]] * 3) else: if uc is not None: if isinstance(pose, list): pose = pose[:N]*2 else: pose = torch.cat([pose[:N]] * 2) sampling_kwargs['pose'] = pose sampling_kwargs['drop_im'] = None sampling_kwargs['mask_ref'] = None for k in c: if isinstance(c[k], torch.Tensor): c[k], uc[k] = map(lambda y: y[k][:(n+1)*N].to('cuda'), (c, uc)) import time st = time.time() with model.ema_scope("Plotting"): samples = model.sample( c, shape=noise.shape[1:], uc=uc, batch_size=N, num_steps=num_steps, noise=noise, **sampling_kwargs ) model.clear_rendered_feat() samples = model.decode_first_stage(samples) print("Time taken for sampling", time.time() - st) log["samples"] = samples.cpu() return log def process_camera_json(camera_json, example_cam): # replace all single quotes in the camera_json with quotes quotes camera_json = camera_json.replace("'", "\"") print("input camera json") print(camera_json) camera_dict = json.loads(camera_json)["scene.camera"] eye = torch.tensor([camera_dict["eye"]["x"], camera_dict["eye"]["y"], camera_dict["eye"]["z"]], dtype=torch.float32).unsqueeze(0) up = torch.tensor([camera_dict["up"]["x"], camera_dict["up"]["y"], camera_dict["up"]["z"]], dtype=torch.float32).unsqueeze(0) center = torch.tensor([camera_dict["center"]["x"], camera_dict["center"]["y"], camera_dict["center"]["z"]], dtype=torch.float32).unsqueeze(0) new_R, new_T = look_at_view_transform(eye=eye, at=center, up=up) ## temp # new_R = torch.tensor([[[ 0.4988, 0.2666, 0.8247], # [-0.1917, -0.8940, 0.4049], # [ 0.8453, -0.3601, -0.3948]]], dtype=torch.float32) # new_T = torch.tensor([[ 0.0739, -0.0013, 0.9973]], dtype=torch.float32) # new_R = torch.tensor([[[ 0.2530, 0.2989, 0.9201], # [-0.2652, -0.8932, 0.3631], # [ 0.9304, -0.3359, -0.1467],]], dtype=torch.float32) # new_T = torch.tensor([[ 0.0081, 0.0337, 1.0452]], dtype=torch.float32) print("focal length", example_cam.focal_length) print("principal point", example_cam.principal_point) newcam = PerspectiveCameras(R=new_R, T=new_T, focal_length=example_cam.focal_length, principal_point=example_cam.principal_point, image_size=512) print("input pose") print(newcam.get_world_to_view_transform().get_matrix()) return newcam def load_and_return_model_and_data(config, model, ckpt="/data/gdsu/customization3d/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors", delta_ckpt=None, train=False, valid=False, far=3, num_images=1, num_ref=8, max_images=20, ): config = OmegaConf.load(config) # load data data = None # config.data.params.jitter = False # config.data.params.addreg = False # config.data.params.bbox = False # data = instantiate_from_config(config.data) # data = data.train_dataset # single_id = data.single_id # if hasattr(data, 'rotations'): # total_images = len(data.rotations[data.sequence_list[single_id]]) # else: # total_images = len(data.annotations['chair']) # print(f"Total images in dataset: {total_images}") model, msg = load_delta_model(model, delta_ckpt,) # change forward methods to store rendered features and use the pre-calculated reference features def register_recr(net_): if net_.__class__.__name__ == 'SpatialTransformer': print(net_.__class__.__name__, "adding control") bound_method = customforward.__get__(net_, net_.__class__) setattr(net_, 'forward', bound_method) return elif hasattr(net_, 'children'): for net__ in net_.children(): register_recr(net__) return def register_recr2(net_): if net_.__class__.__name__ == 'BasicTransformerBlock': print(net_.__class__.__name__, "adding control") bound_method = _customforward.__get__(net_, net_.__class__) setattr(net_, 'forward', bound_method) return elif hasattr(net_, 'children'): for net__ in net_.children(): register_recr2(net__) return sub_nets = model.model.diffusion_model.named_children() for net in sub_nets: register_recr(net[1]) register_recr2(net[1]) # start sampling model.clear_rendered_feat() return model, data def sample(model, data, num_images=1, prompt="", appendpath="", camera_json=None, train=False, scale=7.5, scale_im=3.5, beta=1.0, num_ref=8, skipreflater=False, num_steps=10, valid=False, max_images=20, seed=42, camera_path="pretrained-models/car0/camera.bin", ): """ Only works with num_images=1 (because of camera_json processing) """ if num_images != 1: print("forcing num_images to be 1") num_images = 1 # set guidance scales model.sampler.guider.scale_im = scale_im model.sampler.guider.scale = scale seed_everything(seed) # load cameras cameras_val, cameras_train = torch.load(camera_path) global choices num_ref = 8 max_diff = len(cameras_train)/num_ref choices = [int(x) for x in torch.linspace(0, len(cameras_train) - max_diff, num_ref)] cameras_train_final = [cameras_train[i] for i in choices] # start sampling model.clear_rendered_feat() if prompt == "": prompt = None noise = torch.randn(1, 4, 64, 64).to('cuda').repeat(num_images, 1, 1, 1) # random sample camera poses pose_ids = np.random.choice(len(cameras_val), num_images, replace=False) print(pose_ids) pose_ids[0] = 21 pose = [cameras_val[i] for i in pose_ids] print("example camera") print(pose[0].R) print(pose[0].T) print(pose[0].focal_length) print(pose[0].principal_point) # prepare batches [if translating then call required functions on the target pose] batches = [] for i in range(num_images): batch = {'pose': [pose[i]] + cameras_train_final, "original_size_as_tuple": torch.tensor([512, 512]).reshape(-1, 2), "target_size_as_tuple": torch.tensor([512, 512]).reshape(-1, 2), "crop_coords_top_left": torch.tensor([0, 0]).reshape(-1, 2), "original_size_as_tuple_ref": torch.tensor([512, 512]).reshape(-1, 2), "target_size_as_tuple_ref": torch.tensor([512, 512]).reshape(-1, 2), "crop_coords_top_left_ref": torch.tensor([0, 0]).reshape(-1, 2), } batch_ = copy.deepcopy(batch) batch_["pose"][0] = process_camera_json(camera_json, pose[0]) batch_["pose"] = [join_cameras_as_batch(batch_["pose"])] # print('batched') # print(batch_["pose"][0].get_world_to_view_transform().get_matrix()) batches.append(batch_) print(f'len batches: {len(batches)}') image = None with torch.no_grad(): for batch in batches: for key in batch.keys(): if isinstance(batch[key], torch.Tensor): batch[key] = batch[key].to('cuda') elif 'pose' in key: batch[key] = [x.to('cuda') for x in batch[key]] else: pass if prompt is not None: batch["txt"] = [prompt for _ in range(1)] batch["txt_ref"] = [prompt for _ in range(len(batch["pose"])-1)] print(batch["txt"]) N = 1 log_ = log_images(model, batch, N=N, noise=noise.clone()[:N], num_steps=num_steps, scale_im=scale_im) image = log_["samples"] torch.cuda.empty_cache() model.clear_rendered_feat() print("generation done") return image