forked from diffusion-classifier/diffusion-classifier
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmodels.py
More file actions
74 lines (67 loc) · 2.52 KB
/
models.py
File metadata and controls
74 lines (67 loc) · 2.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler, StableDiffusionPipeline, \
EulerDiscreteScheduler
device = 'cuda' if torch.cuda.is_available() else 'cpu'
MODEL_IDS = {
'1-1': "CompVis/stable-diffusion-v1-1",
'1-2': "CompVis/stable-diffusion-v1-2",
'1-3': "CompVis/stable-diffusion-v1-3",
'1-4': "CompVis/stable-diffusion-v1-4",
'1-5': "runwayml/stable-diffusion-v1-5",
'2-0': "stabilityai/stable-diffusion-2-base",
'2-1': "stabilityai/stable-diffusion-2-1-base",
'small': "OFA-Sys/small-stable-diffusion-v0"
}
def get_sd_model(args):
if args.dtype == 'float32':
dtype = torch.float32
elif args.dtype == 'float16':
dtype = torch.float16
else:
raise NotImplementedError
assert args.version in MODEL_IDS.keys()
model_id = MODEL_IDS[args.version]
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=dtype).to(device)
pipe.enable_xformers_memory_efficient_attention()
vae = pipe.vae
tokenizer = pipe.tokenizer
text_encoder = pipe.text_encoder
# unet = torch.compile(pipe.unet)
unet = pipe.unet
# unet = unet.to(memory_format=torch.channels_last)
return vae, tokenizer, text_encoder, unet, scheduler
def get_scheduler_config(args):
if args.version in {'1-1', '1-2', '1-3', '1-4', '1-5'}:
config = {
"_class_name": "EulerDiscreteScheduler",
"_diffusers_version": "0.14.0",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"interpolation_type": "linear",
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"set_alpha_to_one": False,
"skip_prk_steps": True,
"steps_offset": 1,
"trained_betas": None
}
elif args.version in {'2-0', '2-1'}:
config = {
"_class_name": "EulerDiscreteScheduler",
"_diffusers_version": "0.10.2",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": False,
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"set_alpha_to_one": False,
"skip_prk_steps": True,
"steps_offset": 1, # todo
"trained_betas": None
}
else:
raise NotImplementedError
return config