mirror of
https://github.com/nod-ai/AMD-SHARK-Studio.git
synced 2026-04-03 03:00:17 -04:00
4505c4549f
* Add igpu and custom triple support. * Small fixes to igpu, SDXL-turbo * custom pipe loading * formatting * Remove old nodlogo import.
119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
# from shark_turbine.turbine_models.schedulers import export_scheduler_model
|
|
from diffusers import (
|
|
LCMScheduler,
|
|
LMSDiscreteScheduler,
|
|
PNDMScheduler,
|
|
DDPMScheduler,
|
|
DDIMScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
KDPM2DiscreteScheduler,
|
|
EulerDiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
DEISMultistepScheduler,
|
|
DPMSolverSinglestepScheduler,
|
|
KDPM2AncestralDiscreteScheduler,
|
|
HeunDiscreteScheduler,
|
|
)
|
|
|
|
|
|
def get_schedulers(model_id):
|
|
# TODO: switch over to turbine and run all on GPU
|
|
print(f"\n[LOG] Initializing schedulers from model id: {model_id}")
|
|
schedulers = dict()
|
|
schedulers["PNDM"] = PNDMScheduler.from_pretrained(
|
|
model_id,
|
|
subfolder="scheduler",
|
|
)
|
|
# schedulers["DDPM"] = DDPMScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["KDPM2Discrete"] = KDPM2DiscreteScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["LMSDiscrete"] = LMSDiscreteScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["DDIM"] = DDIMScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["LCMScheduler"] = LCMScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["DPMSolverMultistep"] = DPMSolverMultistepScheduler.from_pretrained(
|
|
# model_id, subfolder="scheduler", algorithm_type="dpmsolver"
|
|
# )
|
|
# schedulers["DPMSolverMultistep++"] = DPMSolverMultistepScheduler.from_pretrained(
|
|
# model_id, subfolder="scheduler", algorithm_type="dpmsolver++"
|
|
# )
|
|
# schedulers["DPMSolverMultistepKarras"] = (
|
|
# DPMSolverMultistepScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# use_karras_sigmas=True,
|
|
# )
|
|
# )
|
|
# schedulers["DPMSolverMultistepKarras++"] = (
|
|
# DPMSolverMultistepScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# algorithm_type="dpmsolver++",
|
|
# use_karras_sigmas=True,
|
|
# )
|
|
# )
|
|
schedulers["EulerDiscrete"] = EulerDiscreteScheduler.from_pretrained(
|
|
model_id,
|
|
subfolder="scheduler",
|
|
)
|
|
schedulers["EulerAncestralDiscrete"] = (
|
|
EulerAncestralDiscreteScheduler.from_pretrained(
|
|
model_id,
|
|
subfolder="scheduler",
|
|
)
|
|
)
|
|
# schedulers["DEISMultistep"] = DEISMultistepScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["DPMSolverSinglestep"] = DPMSolverSinglestepScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# schedulers["KDPM2AncestralDiscrete"] = (
|
|
# KDPM2AncestralDiscreteScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
# )
|
|
# schedulers["HeunDiscrete"] = HeunDiscreteScheduler.from_pretrained(
|
|
# model_id,
|
|
# subfolder="scheduler",
|
|
# )
|
|
return schedulers
|
|
|
|
|
|
def export_scheduler_model(model):
|
|
return "None", "None"
|
|
|
|
|
|
scheduler_model_map = {
|
|
"PNDM": export_scheduler_model("PNDMScheduler"),
|
|
# "DPMSolverSDE": export_scheduler_model("DpmSolverSDEScheduler"),
|
|
"EulerDiscrete": export_scheduler_model("EulerDiscreteScheduler"),
|
|
"EulerAncestralDiscrete": export_scheduler_model("EulerAncestralDiscreteScheduler"),
|
|
# "LCM": export_scheduler_model("LCMScheduler"),
|
|
# "LMSDiscrete": export_scheduler_model("LMSDiscreteScheduler"),
|
|
# "DDPM": export_scheduler_model("DDPMScheduler"),
|
|
# "DDIM": export_scheduler_model("DDIMScheduler"),
|
|
# "DPMSolverMultistep": export_scheduler_model("DPMSolverMultistepScheduler"),
|
|
# "KDPM2Discrete": export_scheduler_model("KDPM2DiscreteScheduler"),
|
|
# "DEISMultistep": export_scheduler_model("DEISMultistepScheduler"),
|
|
# "DPMSolverSinglestep": export_scheduler_model("DPMSolverSingleStepScheduler"),
|
|
# "KDPM2AncestralDiscrete": export_scheduler_model("KDPM2AncestralDiscreteScheduler"),
|
|
# "HeunDiscrete": export_scheduler_model("HeunDiscreteScheduler"),
|
|
}
|