mirror of https://github.com/hpcaitech/ColossalAI
22 lines
502 B
Python
22 lines
502 B
Python
|
'''
|
||
|
|
torchrun --standalone --nproc_per_node=1 debug.py
|
||
|
|
'''
|
||
|
|
|
||
|
|
from diffusers import AutoencoderKL
|
||
|
|
|
||
|
|
import colossalai
|
||
|
|
from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx
|
||
|
|
|
||
|
|
path = "/data/scratch/diffuser/stable-diffusion-v1-4"
|
||
|
|
|
||
|
|
colossalai.launch_from_torch(config={})
|
||
|
|
with ColoInitContext(device='cpu'):
|
||
|
|
vae = AutoencoderKL.from_pretrained(
|
||
|
|
path,
|
||
|
|
subfolder="vae",
|
||
|
|
revision=None,
|
||
|
|
)
|
||
|
|
|
||
|
|
for n, p in vae.named_parameters():
|
||
|
|
print(n)
|