Spaces:
Runtime error
Runtime error
| import torch | |
| from omegaconf import OmegaConf | |
| from ldm.util import instantiate_from_config | |
| from ldm.models.diffusion.ddpm import LatentDiffusion, DDIMSampler | |
| import numpy as np | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| import json | |
| import os | |
| import time | |
| DEBUG = False | |
| def load_model_from_config(config_path, model_name, device='cuda', load=True): | |
| # Load the config file | |
| config = OmegaConf.load(config_path) | |
| # Instantiate the model | |
| model = instantiate_from_config(config.model) | |
| # Download the model file from Hugging Face | |
| if load: | |
| model_file = hf_hub_download(repo_id=model_name, filename="model.safetensors", token=os.getenv('HF_TOKEN')) | |
| print(f"Loading model from {model_name}") | |
| # Load the state dict | |
| state_dict = torch.load(model_file, map_location='cpu') | |
| model.load_state_dict(state_dict, strict=False) | |
| model.to(device) | |
| model.eval() | |
| return model | |
| def sample_frame(model: LatentDiffusion, prompt: str, image_sequence: torch.Tensor, pos_maps=None, leftclick_maps=None): | |
| sampler = DDIMSampler(model) | |
| with torch.no_grad(): | |
| #u_dict = {'c_crossattn': "", 'c_concat': image_sequence} | |
| #uc = model.get_learned_conditioning(u_dict) | |
| #uc = model.enc_concat_seq(uc, u_dict, 'c_concat') | |
| #c_dict = {'c_crossattn': prompt, 'c_concat': image_sequence} | |
| model.eval() | |
| #c = model.get_learned_conditioning(c_dict) | |
| #print (c['c_crossattn'].shape) | |
| #print (c['c_crossattn'][0]) | |
| print (prompt) | |
| # reshape(B, L * C, H, W) | |
| #height, width, channels = image_sequence.shape | |
| # use einsum to reshape | |
| image_sequence = torch.einsum('hwc->chw', image_sequence).unsqueeze(0) | |
| c = {'c_concat': image_sequence} | |
| print (image_sequence.shape, c['c_concat'].shape) | |
| #c = model.enc_concat_seq(c, c_dict, 'c_concat') | |
| # Zero out the corresponding subtensors in c_concat for padding images | |
| #padding_mask = torch.isclose(image_sequence, torch.tensor(-1.0), rtol=1e-5, atol=1e-5).all(dim=(1, 2, 3)).unsqueeze(0) | |
| #print (padding_mask) | |
| #padding_mask = padding_mask.repeat(1, 4) # Repeat mask 4 times for each projected channel | |
| #print (image_sequence.shape, padding_mask.shape, c['c_concat'].shape) | |
| #c['c_concat'] = c['c_concat'] * (~padding_mask.unsqueeze(-1).unsqueeze(-1)) # Zero out the corresponding features | |
| if pos_maps is not None: | |
| pos_map = pos_maps[0] | |
| leftclick_map = torch.cat(leftclick_maps, dim=0) | |
| print (pos_maps[0].shape, c['c_concat'].shape, leftclick_map.shape) | |
| if False and DEBUG: | |
| c['c_concat'] = c['c_concat']*0 | |
| c['c_concat'] = torch.cat([c['c_concat'][:, :, :, :], pos_maps[0].to(c['c_concat'].device).unsqueeze(0), leftclick_map.to(c['c_concat'].device).unsqueeze(0)], dim=1) | |
| print ('sleeping') | |
| #time.sleep(120) | |
| print ('finished sleeping') | |
| DDPM = False | |
| DDPM = True | |
| DDPM = False | |
| if DEBUG: | |
| #c['c_concat'] = c['c_concat']*0 | |
| print ('utils prompt', prompt, c['c_concat'].shape, c.keys()) | |
| print (c['c_concat'].nonzero()) | |
| #print (c['c_concat'][0, 0, :, :]) | |
| if DDPM: | |
| samples_ddim = model.p_sample_loop(cond=c, shape=[1, 4, 48, 64], return_intermediates=False, verbose=True) | |
| else: | |
| samples_ddim, _ = sampler.sample(S=16, | |
| conditioning=c, | |
| batch_size=1, | |
| shape=[4, 48, 64], | |
| verbose=False) | |
| # unconditional_guidance_scale=5.0, | |
| # unconditional_conditioning=uc, | |
| # eta=0) | |
| print ('dfsf1') | |
| if False and DEBUG: | |
| print ('samples_ddim.shape', samples_ddim.shape) | |
| x_samples_ddim = samples_ddim[:, :3] | |
| # upsample to 512 x 384 | |
| x_samples_ddim = torch.nn.functional.interpolate(x_samples_ddim, size=(384, 512), mode='bilinear') | |
| # create a 512 x 384 image and paste the samples_ddim into the center | |
| #x_samples_ddim = torch.zeros((1, 3, 384, 512)) | |
| #x_samples_ddim[:, :, 128:128+48, 160:160+64] = samples_ddim[:, :3] | |
| else: | |
| print ('dfsf2') | |
| data_mean = -0.54 | |
| data_std = 6.78 | |
| data_min = -27.681446075439453 | |
| data_max = 30.854148864746094 | |
| x_samples_ddim = samples_ddim | |
| x_samples_ddim_feedback = x_samples_ddim | |
| x_samples_ddim = x_samples_ddim * data_std + data_mean | |
| x_samples_ddim = model.decode_first_stage(x_samples_ddim) | |
| print ('dfsf3') | |
| #x_samples_ddim = pos_map.to(c['c_concat'].device).unsqueeze(0).expand(-1, 3, -1, -1) | |
| #x_samples_ddim = model.decode_first_stage(x_samples_ddim) | |
| #x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) | |
| x_samples_ddim = torch.clamp(x_samples_ddim, min=-1.0, max=1.0) | |
| return x_samples_ddim.squeeze(0).cpu().numpy(), x_samples_ddim_feedback.squeeze(0) | |
| # Global variables for model and device | |
| #model = None | |
| #device = None | |
| def initialize_model(config_path, model_name): | |
| #global model, device | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = load_model_from_config(config_path, model_name, device) | |
| return model |