Upload server.py with huggingface_hub

server.py

@@ -22,15 +22,15 @@ def load_model():
     vae = AutoencoderKLWan.from_pretrained(
         "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
         subfolder="vae",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.bfloat16,
     )
     pipe = WanPipeline.from_pretrained(
         "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
         vae=vae,
         torch_dtype=torch.bfloat16,
     )
-    pipe.enable_model_cpu_offload()
-    print("[video] Model loaded with CPU offload!", flush=True)
+    pipe.to("cuda")
+    print("[video] Model loaded on CUDA!", flush=True)
 
 @app.route("/health", methods=["GET"])
 def health():