File size: 629 Bytes
80e6c74 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from diffusers import AutoencoderDC, SanaTransformer2DModel
import torch
def build_sana(vision_tower_cfg, **kwargs):
print("="*20, "Building Sana Transformer", vision_tower_cfg.diffusion_name_or_path, "="*20)
sana = SanaTransformer2DModel.from_pretrained(vision_tower_cfg.diffusion_name_or_path, subfolder="transformer", low_cpu_mem_usage=False, ignore_mismatched_sizes=True, torch_dtype=torch.float16)
return sana
def build_vae(vision_tower_cfg, **kwargs):
vae = AutoencoderDC.from_pretrained(vision_tower_cfg.diffusion_name_or_path, subfolder="vae", torch_dtype=torch.float16)
return vae
|