Update custom_pipeline/sde_ve_pipeline.py

custom_pipeline/sde_ve_pipeline.py

@@ -1,5 +1,3 @@
-from ..scheduler import ScoreSdeVeScheduler
-from ..unet import UNet2DModel
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 
@@ -18,10 +16,10 @@ class ScoreSdeVePipeline(DiffusionPipeline):
             A `ScoreSdeVeScheduler` to be used in combination with `unet` to denoise the encoded image.
     """
 
-    unet: UNet2DModel
-    scheduler: ScoreSdeVeScheduler
+    unet
+    scheduler
 
-    def __init__(self, unet: UNet2DModel, scheduler: ScoreSdeVeScheduler):
+    def __init__(self, unet, scheduler):
         super().__init__()
         self.register_modules(unet=unet, scheduler=scheduler)