Update handler.py

handler.py

@@ -31,26 +31,27 @@ class EndpointHandler:
             timestep_spacing="trailing"
         )
 
-        # Initialize video-to-video pipeline
-        self.pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained(
-            path or "jbilcke-hf/CogVideoX-Fun-V1.5-5b-for-InferenceEndpoints",
-            transformer=self.pipe.transformer,
-            vae=self.pipe.vae,
-            scheduler=self.pipe.scheduler,
-            tokenizer=self.pipe.tokenizer,
-            text_encoder=self.pipe.text_encoder,
-            torch_dtype=torch.bfloat16
-        ).to("cuda")
-
-        # Initialize image-to-video pipeline
-        self.pipe_image = CogVideoXImageToVideoPipeline.from_pretrained(
-            path or "THUDM/CogVideoX1.5-5B-I2V",
-            vae=self.pipe.vae,
-            scheduler=self.pipe.scheduler,
-            tokenizer=self.pipe.tokenizer,
-            text_encoder=self.pipe.text_encoder,
-            torch_dtype=torch.bfloat16
-        ).to("cuda")
+        # those two pipelines - generated by Claude - are interesting, but loading it all at once is too much.
+        # # Initialize video-to-video pipeline
+        # self.pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained(
+        #     path or "jbilcke-hf/CogVideoX-Fun-V1.5-5b-for-InferenceEndpoints",
+        #     transformer=self.pipe.transformer,
+        #     vae=self.pipe.vae,
+        #     scheduler=self.pipe.scheduler,
+        #     tokenizer=self.pipe.tokenizer,
+        #     text_encoder=self.pipe.text_encoder,
+        #     torch_dtype=torch.bfloat16
+        # ).to("cuda")
+        # 
+        # # Initialize image-to-video pipeline
+        # self.pipe_image = CogVideoXImageToVideoPipeline.from_pretrained(
+        #     path or "THUDM/CogVideoX1.5-5B-I2V",
+        #     vae=self.pipe.vae,
+        #     scheduler=self.pipe.scheduler,
+        #     tokenizer=self.pipe.tokenizer,
+        #     text_encoder=self.pipe.text_encoder,
+        #     torch_dtype=torch.bfloat16
+        # ).to("cuda")
 
     def _decode_base64_to_image(self, base64_string: str) -> Image.Image:
         """Convert base64 string to PIL Image."""
@@ -101,16 +102,19 @@ class EndpointHandler:
             input_image = self._decode_base64_to_image(data["image"])
             input_image = input_image.resize((720, 480))  # Resize as per example
             image = load_image(input_image)
-            video_frames = self.pipe_image(
-                image=image,
-                **generation_kwargs
-            ).frames[0]
+
+            #raise ValueError("image to video isn't supported yet (takes up too much RAM right now)")
+            return {"error": "Image to video generation not yet supported"}
+            #video_frames = self.pipe_image(
+            #    image=image,
+            #    **generation_kwargs
+            #).frames[0]
         
         elif "video" in data:
             # Video to video generation
             # TODO: Implement video loading from base64
             # For now, returning error
-            return {"error": "Video to video generation not yet implemented"}
+            return {"error": "Video to video generation not yet supported"}
         
         else:
             # Text to video generation
@@ -128,7 +132,7 @@ class EndpointHandler:
         """Cleanup the model and free GPU memory."""
         # Move models to CPU to free GPU memory
         self.pipe.to("cpu")
-        self.pipe_video.to("cpu")
-        self.pipe_image.to("cpu")
+        #self.pipe_video.to("cpu")
+        #self.pipe_image.to("cpu")
         # Clear CUDA cache
         torch.cuda.empty_cache()