File size: 3,577 Bytes
3d1c0e1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
# Copyright (c) 2025 FoundationVision
# SPDX-License-Identifier: MIT
import os
import os.path as osp
import cv2
import torch
import torch.nn.functional as F
import numpy as np
from infinity.schedules.dynamic_resolution import get_first_full_spatial_size_scale_index
def labels2image(all_indices, label_type='int_label', scale_schedule=None):
summed_codes, recons_imgs = self.vae.decode_from_indices(all_indices, scale_schedule, label_type)
recons_img = recons_imgs[0]
recons_img = (recons_img + 1) / 2
recons_img = recons_img.permute(1, 2, 0).mul_(255).cpu().numpy().astype(np.uint8)[:,:,::-1]
return recons_img
def features2image(raw_features):
recons_imgs = self.vae.decode(raw_features.squeeze(-3))
recons_img = recons_imgs[0]
recons_img = (recons_img + 1) / 2
recons_img = recons_img.permute(1, 2, 0).mul_(255).cpu().numpy().astype(np.uint8)[:,:,::-1]
return recons_img
class SelfCorrection(object):
def __init__(self, vae, args):
self.noise_apply_layers = args.noise_apply_layers
self.noise_apply_requant = args.noise_apply_requant
self.noise_apply_strength = args.noise_apply_strength
if not isinstance(self.noise_apply_strength, list):
self.noise_apply_strength = str(self.noise_apply_strength)
self.noise_apply_strength = list(map(float, self.noise_apply_strength.split(',')))
if len(self.noise_apply_strength) == 1:
self.noise_apply_strength = self.noise_apply_strength[0]
self.apply_spatial_patchify = args.apply_spatial_patchify
self.vae = vae
print(f'self.noise_apply_strength: {self.noise_apply_strength}')
def apply_noise_requant(self, bit_indices, quantized, args, device, si, lfq=None, noise_apply_strength=None):
if lfq is None:
lfq = self.vae.quantizer.lfq
if noise_apply_strength is None:
noise_apply_strength = self.noise_apply_strength
if isinstance(noise_apply_strength, list):
noise_apply_strength = np.random.randint(0, max(1, 100 * noise_apply_strength[si]+1)) * 0.01
else:
noise_apply_strength = np.random.randint(0, max(1, 100 * noise_apply_strength+1)) * 0.01
mask = torch.rand(*bit_indices.shape, device=device) < noise_apply_strength
pred_bit_indices = bit_indices.clone()
if args.num_of_label_value == 2:
pred_bit_indices[mask] = 1 - pred_bit_indices[mask]
else:
noise = torch.randint(0, args.num_of_label_value, bit_indices.shape, dtype=bit_indices.dtype, device=device)
pred_bit_indices[mask] = noise[mask]
if self.noise_apply_requant:
quantized = lfq.indices_to_codes(pred_bit_indices, label_type = 'bit_label')
return pred_bit_indices, quantized
def visualize(self, vae_scale_schedule, inp_B3HW, gt_all_bit_indices, pred_all_bit_indices):
gt_img = (inp_B3HW.squeeze(-3) + 1) / 2 * 255
gt_img = gt_img[0].permute(1,2,0).cpu().numpy().astype(np.uint8)[:,:,::-1]
recons_img_2 = labels2image(gt_all_bit_indices, label_type='bit_label', scale_schedule=vae_scale_schedule)
recons_img_3 = labels2image(pred_all_bit_indices, label_type='bit_label', scale_schedule=vae_scale_schedule)
cat_image = np.concatenate([gt_img, recons_img_2, recons_img_3], axis=1)
save_path = osp.abspath('non_teacher_force.jpg')
cv2.imwrite(save_path, cat_image)
print(f'Save to {save_path}')
import pdb; pdb.set_trace()
print(cat_image.shape)
|