red-kit
/

loop-gpt-video-handler

redkits commited on Aug 15

Commit

6e9d2f8

verified ·

1 Parent(s): 8633cca

Upload server.py with huggingface_hub

Files changed (1) hide show

server.py CHANGED Viewed

@@ -29,8 +29,8 @@ def load_model():
         vae=vae,
         torch_dtype=torch.bfloat16,
     )
-    pipe.to("cuda")
-    print("[video] Model loaded successfully!", flush=True)
 @app.route("/health", methods=["GET"])
 def health():

         vae=vae,
         torch_dtype=torch.bfloat16,
     )
+    pipe.enable_model_cpu_offload()
+    print("[video] Model loaded with CPU offload!", flush=True)
 @app.route("/health", methods=["GET"])
 def health():