Update app.py

app.py

@@ -16,6 +16,8 @@ from diffusers.utils import export_to_video
 from diffusers.utils import load_image
 from PIL import Image
 
+from torchao.quantization import float8_weight_only
+from torchao.quantization import quantize_
 from transformers import LlamaModel
 
 import torch
@@ -55,13 +57,27 @@ def init_predictor():
         torch_dtype=torch.bfloat16,
         device="cpu",
     ).to("cpu").eval()
+    quantize_(text_encoder, float8_weight_only(), device="cpu")
+    text_encoder.to("cpu")
+    torch.cuda.empty_cache()
+    quantize_(transformer, float8_weight_only(), device="cpu")
+    transformer.to("cpu")
+    torch.cuda.empty_cache()
+    offload_config=OffloadConfig(
+            high_cpu_memory=True,
+            parameters_level=True,
+            compiler_transformer=False,
+    )
     pipe = SkyreelsVideoPipeline.from_pretrained(
         base_model_id,
         transformer=transformer,
         text_encoder=text_encoder,
         torch_dtype=torch.bfloat16,
     ).to("cpu")
-    pipe.to(device)
+    Offload.offload(
+            pipeline=pipe,
+            config=offload_config,
+    )
     
 negative_prompt = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion"