fix: sdxl build oom (#151)

Author: varshith15

src/streamdiffusion/acceleration/tensorrt/builder.py

@@ -60,6 +60,7 @@ def build(
                 opt_batch_size=opt_batch_size,
                 onnx_opset=onnx_opset,
             )
+            self.network = self.network.to("cpu")
             del self.network
             gc.collect()
             torch.cuda.empty_cache()
@@ -89,7 +90,6 @@ def build(
                 build_all_tactics=build_all_tactics,
                 build_enable_refit=build_enable_refit,
             )
-        
         for file in os.listdir(os.path.dirname(engine_path)):
             if file.endswith('.engine'):
                 continue

src/streamdiffusion/modules/controlnet_module.py

@@ -610,7 +610,7 @@ def _load_pytorch_controlnet_model(self, model_id: str, conditioning_channels: O
                     )
                 else:
                     controlnet = ControlNetModel.from_pretrained(model_id, **load_kwargs)
-            controlnet = controlnet.to(device=self.device, dtype=self.dtype)
+            controlnet = controlnet.to(dtype=self.dtype)
             # Track model_id for updater diffing
             try:
                 setattr(controlnet, 'model_id', model_id)

src/streamdiffusion/wrapper.py

@@ -1042,9 +1042,9 @@ def _load_model(
                 traceback.print_exc()
             raise RuntimeError(error_msg)
         else:
-            if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
+            if not compile_engines_only and hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
                 pipe.text_encoder = pipe.text_encoder.to(device=self.device)
-            if hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:
+            if not compile_engines_only and hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:
                 pipe.text_encoder_2 = pipe.text_encoder_2.to(device=self.device)
 
         # If we get here, the model loaded successfully - break out of retry loop
@@ -1571,7 +1571,7 @@ def _load_model(
                 if self.use_safety_checker or safety_checker_engine_exists:
                     if not safety_checker_engine_exists:
                         from transformers import AutoModelForImageClassification
-                        self.safety_checker = AutoModelForImageClassification.from_pretrained(safety_checker_model_id).to("cuda")
+                        self.safety_checker = AutoModelForImageClassification.from_pretrained(safety_checker_model_id)
 
                         safety_checker_model = NSFWDetector(
                             device=self.device,
@@ -1586,7 +1586,7 @@ def _load_model(
                             model_config=safety_checker_model,
                             batch_size=self.batch_size if self.mode == "txt2img" else stream.frame_bff_size,
                             cuda_stream=None,
-                            load_engine=load_engine,
+                            load_engine=False,
                         )
 
                     if load_engine: