Update handler.py

handler.py

@@ -279,21 +279,25 @@ def create_ltx_video_pipeline(
     """Create and configure the LTX video pipeline"""
     # Get the absolute paths for the model components
     current_dir = Path.cwd()
+
+    ckpt_path = "./txv-2b-0.9.6-distilled-04-25.safetensors"
     
     # Get allowed inference steps from config if available
     allowed_inference_steps = None
-    try:
-        # Load allowed inference steps from metadata if available
-        if Path("transformer/config.json").exists():
-            with open("transformer/config.json", "r") as f:
-                config_data = json.load(f)
-                allowed_inference_steps = config_data.get("allowed_inference_steps")
-    except Exception as e:
-        logger.warning(f"Failed to load allowed_inference_steps from config: {e}")
-    
+
+    assert os.path.exists(
+        ckpt_path
+    ), f"Ckpt path provided (--ckpt_path) {ckpt_path} does not exist"
+
+    with safe_open(ckpt_path, framework="pt") as f:
+        metadata = f.metadata()
+        config_str = metadata.get("config")
+        configs = json.loads(config_str)
+        allowed_inference_steps = configs.get("allowed_inference_steps", None)
+
     # Initialize model components
-    vae = CausalVideoAutoencoder.from_pretrained("vae")
-    transformer = Transformer3DModel.from_pretrained("transformer")
+    vae = CausalVideoAutoencoder.from_pretrained(ckpt_path)
+    transformer = Transformer3DModel.from_pretrained(ckpt_path)
     
     # Use constructor if sampler is specified, otherwise use from_pretrained
     if config.sampler:
@@ -301,11 +305,11 @@ def create_ltx_video_pipeline(
             sampler=("Uniform" if config.sampler.lower() == "uniform" else "LinearQuadratic")
         )
     else:
-        scheduler = RectifiedFlowScheduler.from_pretrained("scheduler")
+        scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)
     
-    text_encoder = T5EncoderModel.from_pretrained("text_encoder")
+    text_encoder = T5EncoderModel.from_pretrained("./text_encoder")
     patchifier = SymmetricPatchifier(patch_size=1)
-    tokenizer = T5Tokenizer.from_pretrained("tokenizer")
+    tokenizer = T5Tokenizer.from_pretrained("./tokenizer")
     
     # Move models to the correct device
     vae = vae.to(device)