ALSv/self-forcing

1# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils2# Apache-2.0 License3# By lllyasviel4 5import os6import cv27import json8import random9import glob10import torch11import einops12import numpy as np13import datetime14import torchvision15 16from PIL import Image17 18 19def min_resize(x, m):20    if x.shape[0] < x.shape[1]:21        s0 = m22        s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))23    else:24        s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))25        s1 = m26    new_max = max(s1, s0)27    raw_max = max(x.shape[0], x.shape[1])28    if new_max < raw_max:29        interpolation = cv2.INTER_AREA30    else:31        interpolation = cv2.INTER_LANCZOS432    y = cv2.resize(x, (s1, s0), interpolation=interpolation)33    return y34 35 36def d_resize(x, y):37    H, W, C = y.shape38    new_min = min(H, W)39    raw_min = min(x.shape[0], x.shape[1])40    if new_min < raw_min:41        interpolation = cv2.INTER_AREA42    else:43        interpolation = cv2.INTER_LANCZOS444    y = cv2.resize(x, (W, H), interpolation=interpolation)45    return y46 47 48def resize_and_center_crop(image, target_width, target_height):49    if target_height == image.shape[0] and target_width == image.shape[1]:50        return image51 52    pil_image = Image.fromarray(image)53    original_width, original_height = pil_image.size54    scale_factor = max(target_width / original_width, target_height / original_height)55    resized_width = int(round(original_width * scale_factor))56    resized_height = int(round(original_height * scale_factor))57    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)58    left = (resized_width - target_width) / 259    top = (resized_height - target_height) / 260    right = (resized_width + target_width) / 261    bottom = (resized_height + target_height) / 262    cropped_image = resized_image.crop((left, top, right, bottom))63    return np.array(cropped_image)64 65 66def resize_and_center_crop_pytorch(image, target_width, target_height):67    B, C, H, W = image.shape68 69    if H == target_height and W == target_width:70        return image71 72    scale_factor = max(target_width / W, target_height / H)73    resized_width = int(round(W * scale_factor))74    resized_height = int(round(H * scale_factor))75 76    resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)77 78    top = (resized_height - target_height) // 279    left = (resized_width - target_width) // 280    cropped = resized[:, :, top:top + target_height, left:left + target_width]81 82    return cropped83 84 85def resize_without_crop(image, target_width, target_height):86    if target_height == image.shape[0] and target_width == image.shape[1]:87        return image88 89    pil_image = Image.fromarray(image)90    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)91    return np.array(resized_image)92 93 94def just_crop(image, w, h):95    if h == image.shape[0] and w == image.shape[1]:96        return image97 98    original_height, original_width = image.shape[:2]99    k = min(original_height / h, original_width / w)100    new_width = int(round(w * k))101    new_height = int(round(h * k))102    x_start = (original_width - new_width) // 2103    y_start = (original_height - new_height) // 2104    cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]105    return cropped_image106 107 108def write_to_json(data, file_path):109    temp_file_path = file_path + ".tmp"110    with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:111        json.dump(data, temp_file, indent=4)112    os.replace(temp_file_path, file_path)113    return114 115 116def read_from_json(file_path):117    with open(file_path, 'rt', encoding='utf-8') as file:118        data = json.load(file)119    return data120 121 122def get_active_parameters(m):123    return {k: v for k, v in m.named_parameters() if v.requires_grad}124 125 126def cast_training_params(m, dtype=torch.float32):127    result = {}128    for n, param in m.named_parameters():129        if param.requires_grad:130            param.data = param.to(dtype)131            result[n] = param132    return result133 134 135def separate_lora_AB(parameters, B_patterns=None):136    parameters_normal = {}137    parameters_B = {}138 139    if B_patterns is None:140        B_patterns = ['.lora_B.', '__zero__']141 142    for k, v in parameters.items():143        if any(B_pattern in k for B_pattern in B_patterns):144            parameters_B[k] = v145        else:146            parameters_normal[k] = v147 148    return parameters_normal, parameters_B149 150 151def set_attr_recursive(obj, attr, value):152    attrs = attr.split(".")153    for name in attrs[:-1]:154        obj = getattr(obj, name)155    setattr(obj, attrs[-1], value)156    return157 158 159def print_tensor_list_size(tensors):160    total_size = 0161    total_elements = 0162 163    if isinstance(tensors, dict):164        tensors = tensors.values()165 166    for tensor in tensors:167        total_size += tensor.nelement() * tensor.element_size()168        total_elements += tensor.nelement()169 170    total_size_MB = total_size / (1024 ** 2)171    total_elements_B = total_elements / 1e9172 173    print(f"Total number of tensors: {len(tensors)}")174    print(f"Total size of tensors: {total_size_MB:.2f} MB")175    print(f"Total number of parameters: {total_elements_B:.3f} billion")176    return177 178 179@torch.no_grad()180def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):181    batch_size = a.size(0)182 183    if b is None:184        b = torch.zeros_like(a)185 186    if mask_a is None:187        mask_a = torch.rand(batch_size) < probability_a188 189    mask_a = mask_a.to(a.device)190    mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))191    result = torch.where(mask_a, a, b)192    return result193 194 195@torch.no_grad()196def zero_module(module):197    for p in module.parameters():198        p.detach().zero_()199    return module200 201 202@torch.no_grad()203def supress_lower_channels(m, k, alpha=0.01):204    data = m.weight.data.clone()205 206    assert int(data.shape[1]) >= k207 208    data[:, :k] = data[:, :k] * alpha209    m.weight.data = data.contiguous().clone()210    return m211 212 213def freeze_module(m):214    if not hasattr(m, '_forward_inside_frozen_module'):215        m._forward_inside_frozen_module = m.forward216    m.requires_grad_(False)217    m.forward = torch.no_grad()(m.forward)218    return m219 220 221def get_latest_safetensors(folder_path):222    safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))223 224    if not safetensors_files:225        raise ValueError('No file to resume!')226 227    latest_file = max(safetensors_files, key=os.path.getmtime)228    latest_file = os.path.abspath(os.path.realpath(latest_file))229    return latest_file230 231 232def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):233    tags = tags_str.split(', ')234    tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))235    prompt = ', '.join(tags)236    return prompt237 238 239def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):240    numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)241    if round_to_int:242        numbers = np.round(numbers).astype(int)243    return numbers.tolist()244 245 246def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):247    edges = np.linspace(0, 1, n + 1)248    points = np.random.uniform(edges[:-1], edges[1:])249    numbers = inclusive + (exclusive - inclusive) * points250    if round_to_int:251        numbers = np.round(numbers).astype(int)252    return numbers.tolist()253 254 255def soft_append_bcthw(history, current, overlap=0):256    if overlap <= 0:257        return torch.cat([history, current], dim=2)258 259    assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"260    assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"261 262    weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)263    blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]264    output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)265 266    return output.to(history)267 268 269def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):270    b, c, t, h, w = x.shape271 272    per_row = b273    for p in [6, 5, 4, 3, 2]:274        if b % p == 0:275            per_row = p276            break277 278    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)279    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5280    x = x.detach().cpu().to(torch.uint8)281    x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)282    torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})283    return x284 285 286def save_bcthw_as_png(x, output_filename):287    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)288    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5289    x = x.detach().cpu().to(torch.uint8)290    x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')291    torchvision.io.write_png(x, output_filename)292    return output_filename293 294 295def save_bchw_as_png(x, output_filename):296    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)297    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5298    x = x.detach().cpu().to(torch.uint8)299    x = einops.rearrange(x, 'b c h w -> c h (b w)')300    torchvision.io.write_png(x, output_filename)301    return output_filename302 303 304def add_tensors_with_padding(tensor1, tensor2):305    if tensor1.shape == tensor2.shape:306        return tensor1 + tensor2307 308    shape1 = tensor1.shape309    shape2 = tensor2.shape310 311    new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))312 313    padded_tensor1 = torch.zeros(new_shape)314    padded_tensor2 = torch.zeros(new_shape)315 316    padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1317    padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2318 319    result = padded_tensor1 + padded_tensor2320    return result321 322 323def print_free_mem():324    torch.cuda.empty_cache()325    free_mem, total_mem = torch.cuda.mem_get_info(0)326    free_mem_mb = free_mem / (1024 ** 2)327    total_mem_mb = total_mem / (1024 ** 2)328    print(f"Free memory: {free_mem_mb:.2f} MB")329    print(f"Total memory: {total_mem_mb:.2f} MB")330    return331 332 333def print_gpu_parameters(device, state_dict, log_count=1):334    summary = {"device": device, "keys_count": len(state_dict)}335 336    logged_params = {}337    for i, (key, tensor) in enumerate(state_dict.items()):338        if i >= log_count:339            break340        logged_params[key] = tensor.flatten()[:3].tolist()341 342    summary["params"] = logged_params343 344    print(str(summary))345    return346 347 348def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):349    from PIL import Image, ImageDraw, ImageFont350 351    txt = Image.new("RGB", (width, height), color="white")352    draw = ImageDraw.Draw(txt)353    font = ImageFont.truetype(font_path, size=size)354 355    if text == '':356        return np.array(txt)357 358    # Split text into lines that fit within the image width359    lines = []360    words = text.split()361    current_line = words[0]362 363    for word in words[1:]:364        line_with_word = f"{current_line} {word}"365        if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:366            current_line = line_with_word367        else:368            lines.append(current_line)369            current_line = word370 371    lines.append(current_line)372 373    # Draw the text line by line374    y = 0375    line_height = draw.textbbox((0, 0), "A", font=font)[3]376 377    for line in lines:378        if y + line_height > height:379            break  # stop drawing if the next line will be outside the image380        draw.text((0, y), line, fill="black", font=font)381        y += line_height382 383    return np.array(txt)384 385 386def blue_mark(x):387    x = x.copy()388    c = x[:, :, 2]389    b = cv2.blur(c, (9, 9))390    x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)391    return x392 393 394def green_mark(x):395    x = x.copy()396    x[:, :, 2] = -1397    x[:, :, 0] = -1398    return x399 400 401def frame_mark(x):402    x = x.copy()403    x[:64] = -1404    x[-64:] = -1405    x[:, :8] = 1406    x[:, -8:] = 1407    return x408 409 410@torch.inference_mode()411def pytorch2numpy(imgs):412    results = []413    for x in imgs:414        y = x.movedim(0, -1)415        y = y * 127.5 + 127.5416        y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)417        results.append(y)418    return results419 420 421@torch.inference_mode()422def numpy2pytorch(imgs):423    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0424    h = h.movedim(-1, 1)425    return h426 427 428@torch.no_grad()429def duplicate_prefix_to_suffix(x, count, zero_out=False):430    if zero_out:431        return torch.cat([x, torch.zeros_like(x[:count])], dim=0)432    else:433        return torch.cat([x, x[:count]], dim=0)434 435 436def weighted_mse(a, b, weight):437    return torch.mean(weight.float() * (a.float() - b.float()) ** 2)438 439 440def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):441    x = (x - x_min) / (x_max - x_min)442    x = max(0.0, min(x, 1.0))443    x = x ** sigma444    return y_min + x * (y_max - y_min)445 446 447def expand_to_dims(x, target_dims):448    return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))449 450 451def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):452    if tensor is None:453        return None454 455    first_dim = tensor.shape[0]456 457    if first_dim == batch_size:458        return tensor459 460    if batch_size % first_dim != 0:461        raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")462 463    repeat_times = batch_size // first_dim464 465    return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))466 467 468def dim5(x):469    return expand_to_dims(x, 5)470 471 472def dim4(x):473    return expand_to_dims(x, 4)474 475 476def dim3(x):477    return expand_to_dims(x, 3)478 479 480def crop_or_pad_yield_mask(x, length):481    B, F, C = x.shape482    device = x.device483    dtype = x.dtype484 485    if F < length:486        y = torch.zeros((B, length, C), dtype=dtype, device=device)487        mask = torch.zeros((B, length), dtype=torch.bool, device=device)488        y[:, :F, :] = x489        mask[:, :F] = True490        return y, mask491 492    return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)493 494 495def extend_dim(x, dim, minimal_length, zero_pad=False):496    original_length = int(x.shape[dim])497 498    if original_length >= minimal_length:499        return x500 501    if zero_pad:502        padding_shape = list(x.shape)503        padding_shape[dim] = minimal_length - original_length504        padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)505    else:506        idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)507        last_element = x[idx]508        padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)509 510    return torch.cat([x, padding], dim=dim)511 512 513def lazy_positional_encoding(t, repeats=None):514    if not isinstance(t, list):515        t = [t]516 517    from diffusers.models.embeddings import get_timestep_embedding518 519    te = torch.tensor(t)520    te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)521 522    if repeats is None:523        return te524 525    te = te[:, None, :].expand(-1, repeats, -1)526 527    return te528 529 530def state_dict_offset_merge(A, B, C=None):531    result = {}532    keys = A.keys()533 534    for key in keys:535        A_value = A[key]536        B_value = B[key].to(A_value)537 538        if C is None:539            result[key] = A_value + B_value540        else:541            C_value = C[key].to(A_value)542            result[key] = A_value + B_value - C_value543 544    return result545 546 547def state_dict_weighted_merge(state_dicts, weights):548    if len(state_dicts) != len(weights):549        raise ValueError("Number of state dictionaries must match number of weights")550 551    if not state_dicts:552        return {}553 554    total_weight = sum(weights)555 556    if total_weight == 0:557        raise ValueError("Sum of weights cannot be zero")558 559    normalized_weights = [w / total_weight for w in weights]560 561    keys = state_dicts[0].keys()562    result = {}563 564    for key in keys:565        result[key] = state_dicts[0][key] * normalized_weights[0]566 567        for i in range(1, len(state_dicts)):568            state_dict_value = state_dicts[i][key].to(result[key])569            result[key] += state_dict_value * normalized_weights[i]570 571    return result572 573 574def group_files_by_folder(all_files):575    grouped_files = {}576 577    for file in all_files:578        folder_name = os.path.basename(os.path.dirname(file))579        if folder_name not in grouped_files:580            grouped_files[folder_name] = []581        grouped_files[folder_name].append(file)582 583    list_of_lists = list(grouped_files.values())584    return list_of_lists585 586 587def generate_timestamp():588    now = datetime.datetime.now()589    timestamp = now.strftime('%y%m%d_%H%M%S')590    milliseconds = f"{int(now.microsecond / 1000):03d}"591    random_number = random.randint(0, 9999)592    return f"{timestamp}_{milliseconds}_{random_number}"593 594 595def write_PIL_image_with_png_info(image, metadata, path):596    from PIL.PngImagePlugin import PngInfo597 598    png_info = PngInfo()599    for key, value in metadata.items():600        png_info.add_text(key, value)601 602    image.save(path, "PNG", pnginfo=png_info)603    return image604 605 606def torch_safe_save(content, path):607    torch.save(content, path + '_tmp')608    os.replace(path + '_tmp', path)609    return path610 611 612def move_optimizer_to_device(optimizer, device):613    for state in optimizer.state.values():614        for k, v in state.items():615            if isinstance(v, torch.Tensor):616                state[k] = v.to(device)617