Update handler.py

handler.py

@@ -21,7 +21,7 @@ class EndpointHandler:
         """
         # Initialize pipeline with bfloat16 for optimal performance as recommended in docs
         self.pipe = CogVideoXPipeline.from_pretrained(
-            path or "THUDM/CogVideoX-5b",
+            path or "jbilcke-hf/CogVideoX-Fun-V1.5-5b-for-InferenceEndpoints",
             torch_dtype=torch.bfloat16
         ).to("cuda")
         
@@ -33,7 +33,7 @@ class EndpointHandler:
 
         # Initialize video-to-video pipeline
         self.pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained(
-            path or "THUDM/CogVideoX-5b",
+            path or "jbilcke-hf/CogVideoX-Fun-V1.5-5b-for-InferenceEndpoints",
             transformer=self.pipe.transformer,
             vae=self.pipe.vae,
             scheduler=self.pipe.scheduler,
@@ -44,7 +44,7 @@ class EndpointHandler:
 
         # Initialize image-to-video pipeline
         self.pipe_image = CogVideoXImageToVideoPipeline.from_pretrained(
-            path or "THUDM/CogVideoX-5b-I2V",
+            path or "THUDM/CogVideoX1.5-5B-I2V",
             vae=self.pipe.vae,
             scheduler=self.pipe.scheduler,
             tokenizer=self.pipe.tokenizer,