Unverified Commit 0248541d authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Conversion] Improve conversion script (#1218)

up
parent 5a59f9b7
...@@ -30,6 +30,9 @@ except ImportError: ...@@ -30,6 +30,9 @@ except ImportError:
from diffusers import ( from diffusers import (
AutoencoderKL, AutoencoderKL,
DDIMScheduler, DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
LDMTextToImagePipeline, LDMTextToImagePipeline,
LMSDiscreteScheduler, LMSDiscreteScheduler,
PNDMScheduler, PNDMScheduler,
...@@ -647,7 +650,7 @@ if __name__ == "__main__": ...@@ -647,7 +650,7 @@ if __name__ == "__main__":
"--scheduler_type", "--scheduler_type",
default="pndm", default="pndm",
type=str, type=str,
help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim']", help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancest', 'dpm']",
) )
parser.add_argument( parser.add_argument(
"--extract_ema", "--extract_ema",
...@@ -686,6 +689,16 @@ if __name__ == "__main__": ...@@ -686,6 +689,16 @@ if __name__ == "__main__":
) )
elif args.scheduler_type == "lms": elif args.scheduler_type == "lms":
scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear") scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear")
elif args.scheduler_type == "euler":
scheduler = EulerDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear")
elif args.scheduler_type == "euler-ancestral":
scheduler = EulerAncestralDiscreteScheduler(
beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear"
)
elif args.scheduler_type == "dpm":
scheduler = DPMSolverMultistepScheduler(
beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear"
)
elif args.scheduler_type == "ddim": elif args.scheduler_type == "ddim":
scheduler = DDIMScheduler( scheduler = DDIMScheduler(
beta_start=beta_start, beta_start=beta_start,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment