Spaces:
Runtime error
Runtime error
| import sys, os | |
| import torch | |
| from .vqvae import VQ, VQVAE, DiVAE, VQControlNet | |
| from .scheduling import * | |
| def get_image_tokenizer(tokenizer_id: str, | |
| tokenizers_root: str = './tokenizer_ckpts', | |
| encoder_only: bool = False, | |
| device: str = 'cuda', | |
| verbose: bool = True, | |
| return_None_on_fail: bool = False,): | |
| """ | |
| Load a pretrained image tokenizer from a checkpoint. | |
| Args: | |
| tokenizer_id (str): ID of the tokenizer to load (name of the checkpoint file without ".pth"). | |
| tokenizers_root (str): Path to the directory containing the tokenizer checkpoints. | |
| encoder_only (bool): Set to True to load only the encoder part of the tokenizer. | |
| device (str): Device to load the tokenizer on. | |
| verbose (bool): Set to True to print load_state_dict warning/success messages | |
| return_None_on_fail (bool): Set to True to return None if the tokenizer fails to load (e.g. doesn't exist) | |
| Returns: | |
| model (nn.Module): The loaded tokenizer. | |
| """ | |
| if return_None_on_fail and not os.path.exists(os.path.join(tokenizers_root, f'{tokenizer_id}.pth')): | |
| return None | |
| if verbose: | |
| print(f'Loading tokenizer {tokenizer_id} ... ', end='') | |
| ckpt = torch.load(os.path.join(tokenizers_root, f'{tokenizer_id}.pth'), map_location='cpu') | |
| # Handle renamed arguments | |
| if 'CLIP' in ckpt['args'].domain or 'DINO' in ckpt['args'].domain or 'ImageBind' in ckpt['args'].domain: | |
| ckpt['args'].patch_proj = False | |
| elif 'sam' in ckpt['args'].domain: | |
| ckpt['args'].input_size_min = ckpt['args'].mask_size | |
| ckpt['args'].input_size_max = ckpt['args'].mask_size | |
| ckpt['args'].input_size = ckpt['args'].mask_size | |
| ckpt['args'].quant_type = getattr(ckpt['args'], 'quantizer_type', None) | |
| ckpt['args'].enc_type = getattr(ckpt['args'], 'encoder_type', None) | |
| ckpt['args'].dec_type = getattr(ckpt['args'], 'decoder_type', None) | |
| ckpt['args'].image_size = getattr(ckpt['args'], 'input_size', None) or getattr(ckpt['args'], 'input_size_max', None) | |
| ckpt['args'].image_size_enc = getattr(ckpt['args'], 'input_size_enc', None) | |
| ckpt['args'].image_size_dec = getattr(ckpt['args'], 'input_size_dec', None) | |
| ckpt['args'].image_size_sd = getattr(ckpt['args'], 'input_size_sd', None) | |
| ckpt['args'].ema_decay = getattr(ckpt['args'], 'quantizer_ema_decay', None) | |
| ckpt['args'].enable_xformer = getattr(ckpt['args'], 'use_xformer', None) | |
| if 'cls_emb.weight' in ckpt['model']: | |
| ckpt['args'].n_labels, ckpt['args'].n_channels = n_labels, n_channels = ckpt['model']['cls_emb.weight'].shape | |
| elif 'encoder.linear_in.weight' in ckpt['model']: | |
| ckpt['args'].n_channels = ckpt['model']['encoder.linear_in.weight'].shape[1] | |
| else: | |
| ckpt['args'].n_channels = ckpt['model']['encoder.proj.weight'].shape[1] | |
| ckpt['args'].sync_codebook = False | |
| if encoder_only: | |
| model_type = VQ | |
| ckpt['model'] = {k: v for k, v in ckpt['model'].items() if 'decoder' not in k and 'post_quant_proj' not in k} | |
| else: | |
| # TODO: Add the model type to the checkpoint when training so we can avoid this hackery | |
| if any(['controlnet' in k for k in ckpt['model'].keys()]): | |
| ckpt['args'].model_type = 'VQControlNet' | |
| elif hasattr(ckpt['args'], 'beta_schedule'): | |
| ckpt['args'].model_type = 'DiVAE' | |
| else: | |
| ckpt['args'].model_type = 'VQVAE' | |
| model_type = getattr(sys.modules[__name__], ckpt['args'].model_type) | |
| model = model_type(**vars(ckpt['args'])) | |
| msg = model.load_state_dict(ckpt['model'], strict=False) | |
| if verbose: | |
| print(msg) | |
| return model.to(device).eval(), ckpt['args'] | |