feat: add support for playground 2.5

This commit is contained in:
Manuel Schmid
2024-06-01 22:43:27 +02:00
parent 07c6c89edf
commit fe58df35c9
9 changed files with 181 additions and 29 deletions
+22 -5
View File
@@ -826,16 +826,33 @@ def worker():
if scheduler_name in ['lcm', 'tcd']:
final_scheduler_name = 'sgm_uniform'
if pipeline.final_unet is not None:
pipeline.final_unet = core.opModelSamplingDiscrete.patch(
def patch_discrete(unet):
return core.opModelSamplingDiscrete.patch(
pipeline.final_unet,
sampling=scheduler_name,
zsnr=False)[0]
if pipeline.final_unet is not None:
pipeline.final_unet = patch_discrete(pipeline.final_unet)
if pipeline.final_refiner_unet is not None:
pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch(
pipeline.final_refiner_unet,
pipeline.final_refiner_unet = patch_discrete(pipeline.final_refiner_unet)
print(f'Using {scheduler_name} scheduler.')
elif scheduler_name == 'edm_playground_v2.5':
final_scheduler_name = 'karras'
def patch_edm(unet):
return core.opModelSamplingContinuousEDM.patch(
unet,
sampling=scheduler_name,
zsnr=False)[0]
sigma_max=120.0,
sigma_min=0.002)[0]
if pipeline.final_unet is not None:
pipeline.final_unet = patch_edm(pipeline.final_unet)
if pipeline.final_refiner_unet is not None:
pipeline.final_refiner_unet = patch_edm(pipeline.final_refiner_unet)
print(f'Using {scheduler_name} scheduler.')
async_task.yields.append(['preview', (flags.preparation_step_count, 'Moving model to GPU ...', None)])