Load schedulers using .from_config

model.py

@@ -51,15 +51,13 @@ class Model:
         elif scheduler_type == 'DDIM':
             pipeline = DDIMPipeline.from_pretrained(repo_id,
                                                     use_auth_token=HF_TOKEN)
-            config, _ = DDIMScheduler.extract_init_dict(
-                dict(pipeline.scheduler.config))
-            pipeline.scheduler = DDIMScheduler(**config)
+            pipeline.scheduler = DDIMScheduler.from_config(
+                repo_id, subfolder='scheduler', use_auth_token=HF_TOKEN)
         elif scheduler_type == 'PNDM':
             pipeline = PNDMPipeline.from_pretrained(repo_id,
                                                     use_auth_token=HF_TOKEN)
-            config, _ = PNDMScheduler.extract_init_dict(
-                dict(pipeline.scheduler.config))
-            pipeline.scheduler = PNDMScheduler(**config)
+            pipeline.scheduler = PNDMScheduler.from_config(
+                repo_id, subfolder='scheduler', use_auth_token=HF_TOKEN)
         else:
             raise ValueError
         return pipeline