Spaces:
Runtime error
Runtime error
| import torch | |
| from tqdm import tqdm | |
| import utils | |
| from PIL import Image | |
| import gc | |
| import numpy as np | |
| from .attention import GatedSelfAttentionDense | |
| from .models import process_input_embeddings, torch_device | |
| def encode(model_dict, image, generator): | |
| """ | |
| image should be a PIL object or numpy array with range 0 to 255 | |
| """ | |
| vae, dtype = model_dict.vae, model_dict.dtype | |
| if isinstance(image, Image.Image): | |
| w, h = image.size | |
| assert w % 8 == 0 and h % 8 == 0, f"h ({h}) and w ({w}) should be a multiple of 8" | |
| # w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 | |
| # image = np.array(image.resize((w, h), resample=Image.Resampling.LANCZOS))[None, :] | |
| image = np.array(image) | |
| if isinstance(image, np.ndarray): | |
| assert image.dtype == np.uint8, f"Should have dtype uint8 (dtype: {image.dtype})" | |
| image = image.astype(np.float32) / 255.0 | |
| image = image[None, ...] | |
| image = image.transpose(0, 3, 1, 2) | |
| image = 2.0 * image - 1.0 | |
| image = torch.from_numpy(image) | |
| assert isinstance(image, torch.Tensor), f"type of image: {type(image)}" | |
| image = image.to(device=torch_device, dtype=dtype) | |
| latents = vae.encode(image).latent_dist.sample(generator) | |
| latents = vae.config.scaling_factor * latents | |
| return latents | |
| def decode(vae, latents): | |
| # scale and decode the image latents with vae | |
| scaled_latents = 1 / 0.18215 * latents | |
| with torch.no_grad(): | |
| image = vae.decode(scaled_latents).sample | |
| image = (image / 2 + 0.5).clamp(0, 1) | |
| image = image.detach().cpu().permute(0, 2, 3, 1).numpy() | |
| images = (image * 255).round().astype("uint8") | |
| return images | |
| def generate(model_dict, latents, input_embeddings, num_inference_steps, guidance_scale = 7.5, no_set_timesteps=False, scheduler_key='dpm_scheduler'): | |
| vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict[scheduler_key], model_dict.dtype | |
| text_embeddings, uncond_embeddings, cond_embeddings = input_embeddings | |
| if not no_set_timesteps: | |
| scheduler.set_timesteps(num_inference_steps) | |
| for t in tqdm(scheduler.timesteps): | |
| # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. | |
| latent_model_input = torch.cat([latents] * 2) | |
| latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t) | |
| # predict the noise residual | |
| with torch.no_grad(): | |
| noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample | |
| # perform guidance | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| # compute the previous noisy sample x_t -> x_t-1 | |
| latents = scheduler.step(noise_pred, t, latents).prev_sample | |
| images = decode(vae, latents) | |
| ret = [latents, images] | |
| return tuple(ret) | |
| def gligen_enable_fuser(unet, enabled=True): | |
| for module in unet.modules(): | |
| if isinstance(module, GatedSelfAttentionDense): | |
| module.enabled = enabled | |
| def prepare_gligen_condition(bboxes, phrases, dtype, tokenizer, text_encoder, num_images_per_prompt): | |
| batch_size = len(bboxes) | |
| assert len(phrases) == len(bboxes) | |
| max_objs = 30 | |
| n_objs = min(max([len(bboxes_item) for bboxes_item in bboxes]), max_objs) | |
| boxes = torch.zeros((batch_size, max_objs, 4), device=torch_device, dtype=dtype) | |
| phrase_embeddings = torch.zeros((batch_size, max_objs, 768), device=torch_device, dtype=dtype) | |
| # masks is a 1D tensor deciding which of the enteries to be enabled | |
| masks = torch.zeros((batch_size, max_objs), device=torch_device, dtype=dtype) | |
| if n_objs > 0: | |
| for idx, (bboxes_item, phrases_item) in enumerate(zip(bboxes, phrases)): | |
| # the length of `bboxes_item` could be smaller than `n_objs` because n_objs takes the max of item length | |
| bboxes_item = torch.tensor(bboxes_item[:n_objs]) | |
| boxes[idx, :bboxes_item.shape[0]] = bboxes_item | |
| tokenizer_inputs = tokenizer(phrases_item[:n_objs], padding=True, return_tensors="pt").to(torch_device) | |
| _phrase_embeddings = text_encoder(**tokenizer_inputs).pooler_output | |
| phrase_embeddings[idx, :_phrase_embeddings.shape[0]] = _phrase_embeddings | |
| assert bboxes_item.shape[0] == _phrase_embeddings.shape[0], f"{bboxes_item.shape[0]} != {_phrase_embeddings.shape[0]}" | |
| masks[idx, :bboxes_item.shape[0]] = 1 | |
| # Classifier-free guidance | |
| repeat_times = num_images_per_prompt * 2 | |
| condition_len = batch_size * repeat_times | |
| boxes = boxes.repeat(repeat_times, 1, 1) | |
| phrase_embeddings = phrase_embeddings.repeat(repeat_times, 1, 1) | |
| masks = masks.repeat(repeat_times, 1) | |
| masks[:condition_len // 2] = 0 | |
| # print("shapes:", boxes.shape, phrase_embeddings.shape, masks.shape) | |
| return boxes, phrase_embeddings, masks, condition_len | |
| def generate_gligen(model_dict, latents, input_embeddings, num_inference_steps, bboxes, phrases, num_images_per_prompt=1, gligen_scheduled_sampling_beta: float = 0.3, guidance_scale=7.5, | |
| frozen_steps=20, frozen_mask=None, | |
| return_saved_cross_attn=False, saved_cross_attn_keys=None, return_cond_ca_only=False, return_token_ca_only=None, | |
| offload_cross_attn_to_cpu=False, offload_latents_to_cpu=True, | |
| return_box_vis=False, show_progress=True, save_all_latents=False, scheduler_key='dpm_scheduler', batched_condition=False): | |
| """ | |
| The `bboxes` should be a list, rather than a list of lists (one box per phrase, we can have multiple duplicated phrases). | |
| """ | |
| vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict[scheduler_key], model_dict.dtype | |
| text_embeddings, _, cond_embeddings = process_input_embeddings(input_embeddings) | |
| if latents.dim() == 5: | |
| # latents_all from the input side, different from the latents_all to be saved | |
| latents_all_input = latents | |
| latents = latents[0] | |
| else: | |
| latents_all_input = None | |
| # Just in case that we have in-place ops | |
| latents = latents.clone() | |
| if save_all_latents: | |
| # offload to cpu to save space | |
| if offload_latents_to_cpu: | |
| latents_all = [latents.cpu()] | |
| else: | |
| latents_all = [latents] | |
| scheduler.set_timesteps(num_inference_steps) | |
| if frozen_mask is not None: | |
| frozen_mask = frozen_mask.to(dtype=dtype).clamp(0., 1.) | |
| # 5.1 Prepare GLIGEN variables | |
| if not batched_condition: | |
| # Add batch dimension to bboxes and phrases | |
| bboxes, phrases = [bboxes], [phrases] | |
| boxes, phrase_embeddings, masks, condition_len = prepare_gligen_condition(bboxes, phrases, dtype, tokenizer, text_encoder, num_images_per_prompt) | |
| if return_saved_cross_attn: | |
| saved_attns = [] | |
| main_cross_attention_kwargs = { | |
| 'offload_cross_attn_to_cpu': offload_cross_attn_to_cpu, | |
| 'return_cond_ca_only': return_cond_ca_only, | |
| 'return_token_ca_only': return_token_ca_only, | |
| 'save_keys': saved_cross_attn_keys, | |
| 'gligen': { | |
| 'boxes': boxes, | |
| 'positive_embeddings': phrase_embeddings, | |
| 'masks': masks | |
| } | |
| } | |
| timesteps = scheduler.timesteps | |
| num_grounding_steps = int(gligen_scheduled_sampling_beta * len(timesteps)) | |
| gligen_enable_fuser(unet, True) | |
| for index, t in enumerate(tqdm(timesteps, disable=not show_progress)): | |
| # Scheduled sampling | |
| if index == num_grounding_steps: | |
| gligen_enable_fuser(unet, False) | |
| # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. | |
| latent_model_input = torch.cat([latents] * 2) | |
| latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t) | |
| main_cross_attention_kwargs['save_attn_to_dict'] = {} | |
| # predict the noise residual | |
| noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, | |
| cross_attention_kwargs=main_cross_attention_kwargs).sample | |
| if return_saved_cross_attn: | |
| saved_attns.append(main_cross_attention_kwargs['save_attn_to_dict']) | |
| del main_cross_attention_kwargs['save_attn_to_dict'] | |
| # perform guidance | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| # compute the previous noisy sample x_t -> x_t-1 | |
| latents = scheduler.step(noise_pred, t, latents).prev_sample | |
| if frozen_mask is not None and index < frozen_steps: | |
| latents = latents_all_input[index+1] * frozen_mask + latents * (1. - frozen_mask) | |
| if save_all_latents: | |
| if offload_latents_to_cpu: | |
| latents_all.append(latents.cpu()) | |
| else: | |
| latents_all.append(latents) | |
| # Turn off fuser for typical SD | |
| gligen_enable_fuser(unet, False) | |
| images = decode(vae, latents) | |
| ret = [latents, images] | |
| if return_saved_cross_attn: | |
| ret.append(saved_attns) | |
| if return_box_vis: | |
| pil_images = [utils.draw_box(Image.fromarray(image), bboxes_item, phrases_item) for image, bboxes_item, phrases_item in zip(images, bboxes, phrases)] | |
| ret.append(pil_images) | |
| if save_all_latents: | |
| latents_all = torch.stack(latents_all, dim=0) | |
| ret.append(latents_all) | |
| return tuple(ret) | |