Update pipeline.py

pipeline.py

@@ -33,31 +33,24 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
 
         """
         super().__init__()
-        self.unet = unet
+        self.register_to_config(
+            batch_size=kwargs.get("batch_size", 1),
+            device=kwargs.get("device", "cuda"),
+            guidance_scale=kwargs.get("guidance_scale", 7.5),
+            lift=kwargs.get("lift", 0.0),
+            num_inference_steps=kwargs.get("num_inference_steps", 50),
+            seed=kwargs.get("seed", 42)
+        )
+
+        # Assign model components
         self.vae = vae
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
         self.scheduler = scheduler
+        self.tokenizer = tokenizer
+        self.unet = unet
+        self.text_encoder = text_encoder
 
-        device = "cuda" if torch.cuda.is_available() else "cpu" 
-
-        self.vae.to(device)
-        self.unet.to(device)
-        self.text_encoder.to(device)
-
-        self.register_to_config(
-            vae=vae.__class__.__name__,
-            scheduler=scheduler.__class__.__name__,
-            tokenizer=tokenizer.__class__.__name__,
-            unet=unet.__class__.__name__,
-            text_encoder=text_encoder.__class__.__name__,
-            device=device,
-            batch_size=None,
-            num_inference_steps=None,
-            guidance_scale=None,
-            lift=None,
-            seed=None,
-        )
+        # Move components to device
+        self.to(torch.device(self.config.device))
 
     @torch.no_grad
     def get_batch(self, latents: Callable, nrow: int, ncol: int) -> Callable: