fix memory bug

inference/pipeline.py

@@ -83,11 +83,16 @@ class RealCustomInferencePipeline:
         vision_model_config = unet_config.pop("vision_model_config", None)
         self.vision_model_config = vision_model_config.pop("vision_model_config", None)
 
-        self.unet_model = UNet2DConditionModelDiffusers(**unet_config)
-
-        self.unet_model.eval().to(self.device).to(self.torch_dtype)
-        self.unet_model.load_state_dict(torch.load(unet_checkpoint, map_location=self.device), strict=False)
-        self.unet_model.load_state_dict(torch.load(realcustom_checkpoint, map_location=self.device), strict=False)
+        # self.unet_model = UNet2DConditionModelDiffusers(**unet_config)
+
+        # self.unet_model.eval().to(self.device).to(self.torch_dtype)
+        # self.unet_model.load_state_dict(torch.load(unet_checkpoint, map_location=self.device), strict=False)
+        # self.unet_model.load_state_dict(torch.load(realcustom_checkpoint, map_location=self.device), strict=False)
+        with torch.device("meta"):
+            self.unet_model = UNet2DConditionModelDiffusers(**unet_config)
+        self.unet_model.load_state_dict(torch.load(unet_checkpoint, map_location=self.device), strict=False, assign=True)
+        self.unet_model.load_state_dict(torch.load(realcustom_checkpoint, map_location=self.device), strict=False, assign=True)
+        self.unet_model.eval()
         print("loading unet model finished.")
 
     def _reload_unet_checkpoint(self, unet_checkpoint, realcustom_checkpoint):