AnchoredAI/llm-grounded-diffusion

1import torch2from tqdm import tqdm3import utils4from PIL import Image5import gc6import numpy as np7from .attention import GatedSelfAttentionDense8from .models import process_input_embeddings, torch_device9 10@torch.no_grad()11def encode(model_dict, image, generator):12    """13    image should be a PIL object or numpy array with range 0 to 25514    """15    16    vae, dtype = model_dict.vae, model_dict.dtype17    18    if isinstance(image, Image.Image):19        w, h = image.size20        assert w % 8 == 0 and h % 8 == 0, f"h ({h}) and w ({w}) should be a multiple of 8"21        # w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 822        # image = np.array(image.resize((w, h), resample=Image.Resampling.LANCZOS))[None, :]23        image = np.array(image)24    25    if isinstance(image, np.ndarray):26        assert image.dtype == np.uint8, f"Should have dtype uint8 (dtype: {image.dtype})"27        image = image.astype(np.float32) / 255.028        image = image[None, ...]29        image = image.transpose(0, 3, 1, 2)30        image = 2.0 * image - 1.031        image = torch.from_numpy(image)32    33    assert isinstance(image, torch.Tensor), f"type of image: {type(image)}"34    35    image = image.to(device=torch_device, dtype=dtype)36    latents = vae.encode(image).latent_dist.sample(generator)37    38    latents = vae.config.scaling_factor * latents39 40    return latents41 42@torch.no_grad()43def decode(vae, latents):44    # scale and decode the image latents with vae45    scaled_latents = 1 / 0.18215 * latents46    with torch.no_grad():47        image = vae.decode(scaled_latents).sample48        49    image = (image / 2 + 0.5).clamp(0, 1)50    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()51    images = (image * 255).round().astype("uint8")52    53    return images54 55@torch.no_grad()56def generate(model_dict, latents, input_embeddings, num_inference_steps, guidance_scale = 7.5, no_set_timesteps=False, scheduler_key='dpm_scheduler'):57    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.dtype58    text_embeddings, uncond_embeddings, cond_embeddings = input_embeddings59    60    if not no_set_timesteps:61        scheduler.set_timesteps(num_inference_steps)62 63    for t in tqdm(scheduler.timesteps):64        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.65        latent_model_input = torch.cat([latents] * 2)66 67        latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)68 69        # predict the noise residual70        with torch.no_grad():71            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample72 73        # perform guidance74        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)75        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)76 77        # compute the previous noisy sample x_t -> x_t-178        latents = scheduler.step(noise_pred, t, latents).prev_sample79 80    images = decode(vae, latents)81    82    ret = [latents, images]83 84    return tuple(ret)85 86def gligen_enable_fuser(unet, enabled=True):87    for module in unet.modules():88        if isinstance(module, GatedSelfAttentionDense):89            module.enabled = enabled90 91def prepare_gligen_condition(bboxes, phrases, dtype, tokenizer, text_encoder, num_images_per_prompt):92    batch_size = len(bboxes)93    94    assert len(phrases) == len(bboxes)95    max_objs = 3096    97    n_objs = min(max([len(bboxes_item) for bboxes_item in bboxes]), max_objs)98    boxes = torch.zeros((batch_size, max_objs, 4), device=torch_device, dtype=dtype)99    phrase_embeddings = torch.zeros((batch_size, max_objs, 768), device=torch_device, dtype=dtype)100    # masks is a 1D tensor deciding which of the enteries to be enabled101    masks = torch.zeros((batch_size, max_objs), device=torch_device, dtype=dtype)102    103    if n_objs > 0:104        for idx, (bboxes_item, phrases_item) in enumerate(zip(bboxes, phrases)):105            # the length of `bboxes_item` could be smaller than `n_objs` because n_objs takes the max of item length106            bboxes_item = torch.tensor(bboxes_item[:n_objs])107            boxes[idx, :bboxes_item.shape[0]] = bboxes_item108 109            tokenizer_inputs = tokenizer(phrases_item[:n_objs], padding=True, return_tensors="pt").to(torch_device)110            _phrase_embeddings = text_encoder(**tokenizer_inputs).pooler_output111            phrase_embeddings[idx, :_phrase_embeddings.shape[0]] = _phrase_embeddings112            assert bboxes_item.shape[0] == _phrase_embeddings.shape[0], f"{bboxes_item.shape[0]} != {_phrase_embeddings.shape[0]}"113            114            masks[idx, :bboxes_item.shape[0]] = 1115 116    # Classifier-free guidance117    repeat_times = num_images_per_prompt * 2118    condition_len = batch_size * repeat_times119 120    boxes = boxes.repeat(repeat_times, 1, 1)121    phrase_embeddings = phrase_embeddings.repeat(repeat_times, 1, 1)122    masks = masks.repeat(repeat_times, 1)123    masks[:condition_len // 2] = 0124    125    # print("shapes:", boxes.shape, phrase_embeddings.shape, masks.shape)126    127    return boxes, phrase_embeddings, masks, condition_len128 129@torch.no_grad()130def 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, 131    frozen_steps=20, frozen_mask=None,132    return_saved_cross_attn=False, saved_cross_attn_keys=None, return_cond_ca_only=False, return_token_ca_only=None, 133    offload_cross_attn_to_cpu=False, offload_latents_to_cpu=True,134    return_box_vis=False, show_progress=True, save_all_latents=False, scheduler_key='dpm_scheduler', batched_condition=False):135    """136    The `bboxes` should be a list, rather than a list of lists (one box per phrase, we can have multiple duplicated phrases).137    """138    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.dtype139    140    text_embeddings, _, cond_embeddings = process_input_embeddings(input_embeddings)141    142    if latents.dim() == 5:143        # latents_all from the input side, different from the latents_all to be saved144        latents_all_input = latents145        latents = latents[0]146    else:147        latents_all_input = None148    149    # Just in case that we have in-place ops150    latents = latents.clone()151    152    if save_all_latents:153        # offload to cpu to save space154        if offload_latents_to_cpu:155            latents_all = [latents.cpu()]156        else:157            latents_all = [latents]158    159    scheduler.set_timesteps(num_inference_steps)160    161    if frozen_mask is not None:162        frozen_mask = frozen_mask.to(dtype=dtype).clamp(0., 1.)163 164    # 5.1 Prepare GLIGEN variables165    if not batched_condition:166        # Add batch dimension to bboxes and phrases167        bboxes, phrases = [bboxes], [phrases]168    169    boxes, phrase_embeddings, masks, condition_len = prepare_gligen_condition(bboxes, phrases, dtype, tokenizer, text_encoder, num_images_per_prompt)170    171    if return_saved_cross_attn:172        saved_attns = []173    174    main_cross_attention_kwargs = {175        'offload_cross_attn_to_cpu': offload_cross_attn_to_cpu,176        'return_cond_ca_only': return_cond_ca_only,177        'return_token_ca_only': return_token_ca_only,178        'save_keys': saved_cross_attn_keys,179        'gligen': {180            'boxes': boxes,181            'positive_embeddings': phrase_embeddings,182            'masks': masks183        }184    }185    186    timesteps = scheduler.timesteps187 188    num_grounding_steps = int(gligen_scheduled_sampling_beta * len(timesteps))189    gligen_enable_fuser(unet, True)190 191    for index, t in enumerate(tqdm(timesteps, disable=not show_progress)):192        # Scheduled sampling193        if index == num_grounding_steps:194            gligen_enable_fuser(unet, False)195        196        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.197        latent_model_input = torch.cat([latents] * 2)198 199        latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)200 201        main_cross_attention_kwargs['save_attn_to_dict'] = {}202 203        # predict the noise residual204        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, 205                            cross_attention_kwargs=main_cross_attention_kwargs).sample206        207        if return_saved_cross_attn:208            saved_attns.append(main_cross_attention_kwargs['save_attn_to_dict'])209            210            del main_cross_attention_kwargs['save_attn_to_dict']211 212        # perform guidance213        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)214        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)215 216        # compute the previous noisy sample x_t -> x_t-1217        latents = scheduler.step(noise_pred, t, latents).prev_sample218        219        if frozen_mask is not None and index < frozen_steps:220            latents = latents_all_input[index+1] * frozen_mask + latents * (1. - frozen_mask)221        222        if save_all_latents:223            if offload_latents_to_cpu:224                latents_all.append(latents.cpu())225            else:226                latents_all.append(latents)227 228    # Turn off fuser for typical SD229    gligen_enable_fuser(unet, False)230    images = decode(vae, latents)231    232    ret = [latents, images]233    if return_saved_cross_attn:234        ret.append(saved_attns)235    if return_box_vis:236        pil_images = [utils.draw_box(Image.fromarray(image), bboxes_item, phrases_item) for image, bboxes_item, phrases_item in zip(images, bboxes, phrases)]237        ret.append(pil_images)238    if save_all_latents:239        latents_all = torch.stack(latents_all, dim=0)240        ret.append(latents_all)241    242    return tuple(ret)243 244