diff options
Diffstat (limited to 'infer.py')
-rw-r--r-- | infer.py | 7 |
1 files changed, 5 insertions, 2 deletions
@@ -23,6 +23,7 @@ from diffusers import ( | |||
23 | EulerAncestralDiscreteScheduler, | 23 | EulerAncestralDiscreteScheduler, |
24 | KDPM2DiscreteScheduler, | 24 | KDPM2DiscreteScheduler, |
25 | KDPM2AncestralDiscreteScheduler, | 25 | KDPM2AncestralDiscreteScheduler, |
26 | DEISMultistepScheduler, | ||
26 | UniPCMultistepScheduler | 27 | UniPCMultistepScheduler |
27 | ) | 28 | ) |
28 | from transformers import CLIPTextModel | 29 | from transformers import CLIPTextModel |
@@ -126,13 +127,13 @@ def create_cmd_parser(): | |||
126 | parser.add_argument( | 127 | parser.add_argument( |
127 | "--scheduler", | 128 | "--scheduler", |
128 | type=str, | 129 | type=str, |
129 | choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "unipc"], | 130 | choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "deis", "unipc"], |
130 | ) | 131 | ) |
131 | parser.add_argument( | 132 | parser.add_argument( |
132 | "--subscheduler", | 133 | "--subscheduler", |
133 | type=str, | 134 | type=str, |
134 | default=None, | 135 | default=None, |
135 | choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"], | 136 | choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a", "deis"], |
136 | ) | 137 | ) |
137 | parser.add_argument( | 138 | parser.add_argument( |
138 | "--template", | 139 | "--template", |
@@ -252,6 +253,8 @@ def create_scheduler(config, scheduler: str, subscheduler: Optional[str] = None) | |||
252 | return KDPM2DiscreteScheduler.from_config(config) | 253 | return KDPM2DiscreteScheduler.from_config(config) |
253 | elif scheduler == "kdpm2_a": | 254 | elif scheduler == "kdpm2_a": |
254 | return KDPM2AncestralDiscreteScheduler.from_config(config) | 255 | return KDPM2AncestralDiscreteScheduler.from_config(config) |
256 | elif scheduler == "deis": | ||
257 | return DEISMultistepScheduler.from_config(config) | ||
255 | elif scheduler == "unipc": | 258 | elif scheduler == "unipc": |
256 | if subscheduler is None: | 259 | if subscheduler is None: |
257 | return UniPCMultistepScheduler.from_config(config) | 260 | return UniPCMultistepScheduler.from_config(config) |