Update handler.py

handler.py

@@ -3,20 +3,14 @@ import torch
 
 class EndpointHandler:
     def __init__(self, model_dir: str):
-        # Load the pipeline strictly from your repo files
         self.pipe = DiffusionPipeline.from_pretrained(
             model_dir,
-            custom_pipeline="pipeline_wan_i2v",   # must match your filename exactly
+            custom_pipeline="pipeline_wan_i2v",
             torch_dtype=torch.float16,
-            local_files_only=True                # prevents HF from downloading anything
+            safety_checker=None
         ).to("cuda")
 
     def __call__(self, data):
-        # Expecting the input: {"image": <base64 or URL or HF file>}
         image = data["image"]
-
-        # Run the model
         result = self.pipe(image)
-
-        # The pipeline returns frames of the generated video
         return {"video": result.frames}