new

pipeline/kiss3d_wrapper.py

@@ -91,7 +91,7 @@ def init_wrapper_from_config(config_path):
     # load lora weights
     flux_pipe.load_lora_weights(flux_lora_pth)
     # flux_pipe.to(device=flux_device)
-    flux_pipe.enable_model_cpu_offload(device=flux_device)
+    # flux_pipe.enable_model_cpu_offload(device=flux_device)
     # flux_pipe = None
 
     # load redux model
@@ -490,6 +490,7 @@ class kiss3d_wrapper(object):
         """
         return: gen_3d_bundle_image, torch.Tensor of shape (3, 1024, 2048), range [0., 1.]
         """
+        self.flux_pipeline.to(self.config['flux'].get('device', 'cpu'))
         print(f"==> generate_3d_bundle_image_text: {prompt}")
         if isinstance(self.flux_pipeline, FluxImg2ImgPipeline):
             flux_pipeline = self.flux_pipeline
@@ -548,7 +549,7 @@ class kiss3d_wrapper(object):
             torchvision.utils.save_image(gen_3d_bundle_image_, save_path)
             logger.info(f"Save generated 3D bundle image to {save_path}")
             return gen_3d_bundle_image_, save_path
-
+        self.flux_pipeline.to('cpu')
         return gen_3d_bundle_image_
     
     def reconstruct_3d_bundle_image(self,