| 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 | |