Update pipeline.py

pipeline.py

@@ -4,13 +4,24 @@ from .generate import Generator
 from .utils import init_model, seed_everything, get_frame_ids
 
 class VidToMePipeline(DiffusionPipeline):
-    def __init__(self, device="cuda", sd_version="2.1", float_precision="fp16", height=512, width=512):
-        # this will initlize the core pipeline components
-        pipe, scheduler, model_key = init_model(device, sd_version, None, "none", float_precision)
-        self.pipe = pipe
-        self.scheduler = scheduler
-        self.model_key = model_key
-        self.device = device
+    # def __init__(self, device="cuda", sd_version="2.1", float_precision="fp16", height=512, width=512):
+    #     # this will initlize the core pipeline components
+    #     pipe, scheduler, model_key = init_model(device, sd_version, None, "none", float_precision)
+    #     self.pipe = pipe
+    #     self.scheduler = scheduler
+    #     self.model_key = model_key
+    #     self.device = device
+    #     self.sd_version = sd_version
+    #     self.float_precision = float_precision
+    #     self.height = height
+    #     self.width = width
+
+    def __init__(self, device="cuda", sd_version="1.5", float_precision="fp16", height=512, width=512):
+        # Register configuration parameters
+        self.register_to_config(device=device, sd_version=sd_version, float_precision=float_precision, height=height, width=width)
+        
+        # Now you can safely use self.device and other attributes
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
         self.sd_version = sd_version
         self.float_precision = float_precision
         self.height = height