@@ -33,7 +33,7 @@ class ScoreSdeVePipeline(DiffusionPipeline):
     unet: UNet2DModel
     scheduler: ScoreSdeVeScheduler
 
-    def __init__(self, unet: UNet2DModel, scheduler: DiffusionPipeline):
+    def __init__(self, unet: UNet2DModel, scheduler: ScoreSdeVeScheduler):
         super().__init__()
         self.register_modules(unet=unet, scheduler=scheduler)