(params: SamplingParams)
| 326 | |
| 327 | |
| 328 | def get_sampler_config(params: SamplingParams): |
| 329 | discretization_config = get_discretization_config(params) |
| 330 | guider_config = get_guider_config(params) |
| 331 | sampler = None |
| 332 | if params.sampler == Sampler.EULER_EDM: |
| 333 | return EulerEDMSampler( |
| 334 | num_steps=params.steps, |
| 335 | discretization_config=discretization_config, |
| 336 | guider_config=guider_config, |
| 337 | s_churn=params.s_churn, |
| 338 | s_tmin=params.s_tmin, |
| 339 | s_tmax=params.s_tmax, |
| 340 | s_noise=params.s_noise, |
| 341 | verbose=True, |
| 342 | ) |
| 343 | if params.sampler == Sampler.HEUN_EDM: |
| 344 | return HeunEDMSampler( |
| 345 | num_steps=params.steps, |
| 346 | discretization_config=discretization_config, |
| 347 | guider_config=guider_config, |
| 348 | s_churn=params.s_churn, |
| 349 | s_tmin=params.s_tmin, |
| 350 | s_tmax=params.s_tmax, |
| 351 | s_noise=params.s_noise, |
| 352 | verbose=True, |
| 353 | ) |
| 354 | if params.sampler == Sampler.EULER_ANCESTRAL: |
| 355 | return EulerAncestralSampler( |
| 356 | num_steps=params.steps, |
| 357 | discretization_config=discretization_config, |
| 358 | guider_config=guider_config, |
| 359 | eta=params.eta, |
| 360 | s_noise=params.s_noise, |
| 361 | verbose=True, |
| 362 | ) |
| 363 | if params.sampler == Sampler.DPMPP2S_ANCESTRAL: |
| 364 | return DPMPP2SAncestralSampler( |
| 365 | num_steps=params.steps, |
| 366 | discretization_config=discretization_config, |
| 367 | guider_config=guider_config, |
| 368 | eta=params.eta, |
| 369 | s_noise=params.s_noise, |
| 370 | verbose=True, |
| 371 | ) |
| 372 | if params.sampler == Sampler.DPMPP2M: |
| 373 | return DPMPP2MSampler( |
| 374 | num_steps=params.steps, |
| 375 | discretization_config=discretization_config, |
| 376 | guider_config=guider_config, |
| 377 | verbose=True, |
| 378 | ) |
| 379 | if params.sampler == Sampler.LINEAR_MULTISTEP: |
| 380 | return LinearMultistepSampler( |
| 381 | num_steps=params.steps, |
| 382 | discretization_config=discretization_config, |
| 383 | guider_config=guider_config, |
| 384 | order=params.order, |
| 385 | verbose=True, |
no test coverage detected