Buckets:
| import copy | |
| import math | |
| from diffusers import ( | |
| DDPMScheduler, | |
| EulerAncestralDiscreteScheduler, | |
| DPMSolverMultistepScheduler, | |
| DPMSolverSinglestepScheduler, | |
| LMSDiscreteScheduler, | |
| PNDMScheduler, | |
| DDIMScheduler, | |
| EulerDiscreteScheduler, | |
| HeunDiscreteScheduler, | |
| KDPM2DiscreteScheduler, | |
| KDPM2AncestralDiscreteScheduler, | |
| LCMScheduler, | |
| FlowMatchEulerDiscreteScheduler, | |
| ) | |
| from toolkit.samplers.mean_flow_scheduler import MeanFlowScheduler | |
| from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler | |
| from toolkit.samplers.custom_lcm_scheduler import CustomLCMScheduler | |
| # scheduler: | |
| SCHEDULER_LINEAR_START = 0.00085 | |
| SCHEDULER_LINEAR_END = 0.0120 | |
| SCHEDULER_TIMESTEPS = 1000 | |
| SCHEDLER_SCHEDULE = "scaled_linear" | |
| sd_config = { | |
| "_class_name": "EulerAncestralDiscreteScheduler", | |
| "_diffusers_version": "0.24.0.dev0", | |
| "beta_end": 0.012, | |
| "beta_schedule": "scaled_linear", | |
| "beta_start": 0.00085, | |
| "clip_sample": False, | |
| "interpolation_type": "linear", | |
| "num_train_timesteps": 1000, | |
| "prediction_type": "epsilon", | |
| "sample_max_value": 1.0, | |
| "set_alpha_to_one": False, | |
| # "skip_prk_steps": False, # for training | |
| "skip_prk_steps": True, | |
| # "steps_offset": 1, | |
| "steps_offset": 0, | |
| # "timestep_spacing": "trailing", # for training | |
| "timestep_spacing": "leading", | |
| "trained_betas": None | |
| } | |
| pixart_config = { | |
| "_class_name": "DPMSolverMultistepScheduler", | |
| "_diffusers_version": "0.22.0.dev0", | |
| "algorithm_type": "dpmsolver++", | |
| "beta_end": 0.02, | |
| "beta_schedule": "linear", | |
| "beta_start": 0.0001, | |
| "dynamic_thresholding_ratio": 0.995, | |
| "euler_at_final": False, | |
| # "lambda_min_clipped": -Infinity, | |
| "lambda_min_clipped": -math.inf, | |
| "lower_order_final": True, | |
| "num_train_timesteps": 1000, | |
| "prediction_type": "epsilon", | |
| "sample_max_value": 1.0, | |
| "solver_order": 2, | |
| "solver_type": "midpoint", | |
| "steps_offset": 0, | |
| "thresholding": False, | |
| "timestep_spacing": "linspace", | |
| "trained_betas": None, | |
| "use_karras_sigmas": False, | |
| "use_lu_lambdas": False, | |
| "variance_type": None | |
| } | |
| flux_config = { | |
| "_class_name": "FlowMatchEulerDiscreteScheduler", | |
| "_diffusers_version": "0.30.0.dev0", | |
| "base_image_seq_len": 256, | |
| "base_shift": 0.5, | |
| "max_image_seq_len": 4096, | |
| "max_shift": 1.15, | |
| "num_train_timesteps": 1000, | |
| "shift": 3.0, | |
| "use_dynamic_shifting": True | |
| } | |
| sd_flow_config = { | |
| "_class_name": "FlowMatchEulerDiscreteScheduler", | |
| "_diffusers_version": "0.30.0.dev0", | |
| "base_image_seq_len": 256, | |
| "base_shift": 0.5, | |
| "max_image_seq_len": 4096, | |
| "max_shift": 1.15, | |
| "num_train_timesteps": 1000, | |
| "shift": 3.0, | |
| "use_dynamic_shifting": False | |
| } | |
| lumina2_config = { | |
| "_class_name": "FlowMatchEulerDiscreteScheduler", | |
| "_diffusers_version": "0.33.0.dev0", | |
| "base_image_seq_len": 256, | |
| "base_shift": 0.5, | |
| "invert_sigmas": False, | |
| "max_image_seq_len": 4096, | |
| "max_shift": 1.15, | |
| "num_train_timesteps": 1000, | |
| "shift": 6.0, | |
| "shift_terminal": None, | |
| "use_beta_sigmas": False, | |
| "use_dynamic_shifting": False, | |
| "use_exponential_sigmas": False, | |
| "use_karras_sigmas": False | |
| } | |
| def get_sampler( | |
| sampler: str, | |
| kwargs: dict = None, | |
| arch: str = "sd" | |
| ): | |
| sched_init_args = {} | |
| if kwargs is not None: | |
| sched_init_args.update(kwargs) | |
| config_to_use = copy.deepcopy(sd_config) if arch == "sd" else copy.deepcopy(pixart_config) | |
| if sampler.startswith("k_"): | |
| sched_init_args["use_karras_sigmas"] = True | |
| if sampler == "ddim": | |
| scheduler_cls = DDIMScheduler | |
| elif sampler == "ddpm": # ddpm is not supported ? | |
| scheduler_cls = DDPMScheduler | |
| elif sampler == "pndm": | |
| scheduler_cls = PNDMScheduler | |
| elif sampler == "lms" or sampler == "k_lms": | |
| scheduler_cls = LMSDiscreteScheduler | |
| elif sampler == "euler" or sampler == "k_euler": | |
| scheduler_cls = EulerDiscreteScheduler | |
| elif sampler == "euler_a": | |
| scheduler_cls = EulerAncestralDiscreteScheduler | |
| elif sampler == "dpmsolver" or sampler == "dpmsolver++" or sampler == "k_dpmsolver" or sampler == "k_dpmsolver++": | |
| scheduler_cls = DPMSolverMultistepScheduler | |
| sched_init_args["algorithm_type"] = sampler.replace("k_", "") | |
| elif sampler == "dpmsingle": | |
| scheduler_cls = DPMSolverSinglestepScheduler | |
| elif sampler == "heun": | |
| scheduler_cls = HeunDiscreteScheduler | |
| elif sampler == "dpm_2": | |
| scheduler_cls = KDPM2DiscreteScheduler | |
| elif sampler == "dpm_2_a": | |
| scheduler_cls = KDPM2AncestralDiscreteScheduler | |
| elif sampler == "lcm": | |
| scheduler_cls = LCMScheduler | |
| elif sampler == "custom_lcm": | |
| scheduler_cls = CustomLCMScheduler | |
| elif sampler == "mean_flow": | |
| scheduler_cls = MeanFlowScheduler | |
| elif sampler == "flowmatch": | |
| scheduler_cls = CustomFlowMatchEulerDiscreteScheduler | |
| config_to_use = copy.deepcopy(flux_config) | |
| if arch == "sd": | |
| config_to_use = copy.deepcopy(sd_flow_config) | |
| elif arch == "flux": | |
| config_to_use = copy.deepcopy(flux_config) | |
| elif arch == "lumina2": | |
| config_to_use = copy.deepcopy(lumina2_config) | |
| else: | |
| print(f"Unknown architecture {arch}, using default flux config") | |
| # use flux by default | |
| config_to_use = copy.deepcopy(flux_config) | |
| else: | |
| raise ValueError(f"Sampler {sampler} not supported") | |
| config = copy.deepcopy(config_to_use) | |
| config.update(sched_init_args) | |
| scheduler = scheduler_cls.from_config(config) | |
| return scheduler | |
Xet Storage Details
- Size:
- 5.6 kB
- Xet hash:
- eadb1a106fd5ee51979925cb0adae512c600f2419d0599e749542ced7aef6ef9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.