Fix set_timestamps for new LCMScheduler implementation

2024-06-28 15:52:54 +02:00
parent 21ab05fd4f
commit d4bda0e740
1 changed files with 1 additions and 1 deletions
@@ -63,7 +63,7 @@ def get_encoder_hidden_states_shape(unet_config, batch_size):

 def get_scheduler():
    scheduler = LCMScheduler.from_pretrained(MODEL_VERSION, subfolder="scheduler")
-    scheduler.set_timesteps(50, get_torch_device(), 50, 1.0)
+    scheduler.set_timesteps(50, get_torch_device(), 50)
    return scheduler