Update testlcm/backend/pipelines/lcm.py

testlcm/backend/pipelines/lcm.py

@@ -24,12 +24,18 @@ def _get_lcm_pipeline_from_base_model(
     unet = UNet2DConditionModel.from_pretrained(
         lcm_model_id,
         torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        safety_checker=None,
+        requires_safety_checker=False,
         local_files_only=use_local_model,
     )
     pipeline = DiffusionPipeline.from_pretrained(
         base_model_id,
         unet=unet,
         torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        safety_checker=None,
+        requires_safety_checker=False,
         local_files_only=use_local_model,
     )
     pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
@@ -40,12 +46,18 @@ def load_taesd(
     pipeline: Any,
     use_local_model: bool = False,
     torch_data_type: torch.dtype = torch.float32,
+    low_cpu_mem_usage=True,
+    safety_checker=None,
+    requires_safety_checker=False,
 ):
     vae_model = get_tiny_decoder_vae_model(pipeline.__class__.__name__)
     pipeline.vae = AutoencoderTiny.from_pretrained(
         vae_model,
         torch_dtype=torch_data_type,
         local_files_only=use_local_model,
+        low_cpu_mem_usage=True,
+        safety_checker=None,
+        requires_safety_checker=False,
     )