engrjaved commited on
Commit
3fe3503
·
1 Parent(s): 306187e

offload text_encoder and transformer to CPU

Browse files
Files changed (1) hide show
  1. inference_script.py +5 -0
inference_script.py CHANGED
@@ -468,6 +468,9 @@ def process_video(
468
  if pipe.transformer.config.use_rotary_positional_embeddings
469
  else None
470
  )
 
 
 
471
 
472
  # Predict noise
473
  predicted_noise = pipe.transformer(
@@ -478,6 +481,8 @@ def process_video(
478
  return_dict=False,
479
  )[0]
480
 
 
 
481
 
482
  latent_generate = pipe.scheduler.get_velocity(
483
  predicted_noise, latent, timesteps
 
468
  if pipe.transformer.config.use_rotary_positional_embeddings
469
  else None
470
  )
471
+
472
+ # offload text encoder to CPU
473
+ pipe.text_encoder.to("cpu")
474
 
475
  # Predict noise
476
  predicted_noise = pipe.transformer(
 
481
  return_dict=False,
482
  )[0]
483
 
484
+ # offload transformer to CPU
485
+ pipe.transformer.to("cpu")
486
 
487
  latent_generate = pipe.scheduler.get_velocity(
488
  predicted_noise, latent, timesteps