Update pipeline.py

pipeline.py

@@ -33,24 +33,31 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
 
         """
         super().__init__()
-        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.scheduler = scheduler
-        self.tokenizer = tokenizer
         self.unet = unet
+        self.vae = vae
         self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.scheduler = scheduler
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        # Move components to device
-        self.to(torch.device(self.config.device))
+        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,
+        )
 
     @torch.no_grad
     def get_batch(self, latents: Callable, nrow: int, ncol: int) -> Callable: