engrjaved
commited on
Commit
·
3fe3503
1
Parent(s):
306187e
offload text_encoder and transformer to CPU
Browse files- 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
|