Update app(workingwithsafetensorscache).py

app(workingwithsafetensorscache).py

@@ -36,9 +36,12 @@ import random
 import tempfile
 from pathlib import Path
 import gc
+import json
 import hashlib
 
+import requests
 import torch
+from safetensors import safe_open
 torch._dynamo.config.suppress_errors = True
 torch._dynamo.config.disable = True
 
@@ -264,6 +267,9 @@ class LTX23DistilledA2VPipeline(DistilledPipeline):
 # Model repos
 LTX_MODEL_REPO = "Lightricks/LTX-2.3"
 GEMMA_REPO ="Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
+DISTILLED_LORA_FILE = "ltx-2.3-22b-distilled-lora-384.safetensors"
+EROS_REPO = "dagloop5/LoRA"
+EROS_FILE = "ltx2310eros_beta.safetensors"
 
 
 # Download model checkpoints
@@ -280,16 +286,50 @@ PENDING_LORA_KEY: str | None = None
 PENDING_LORA_STATE: dict[str, torch.Tensor] | None = None
 PENDING_LORA_STATUS: str = "No LoRA state prepared yet."
 
-weights_dir = Path("weights")
-weights_dir.mkdir(exist_ok=True)
-checkpoint_path = hf_hub_download(
-    repo_id=LTX_MODEL_REPO,
-    filename="ltx-2.3-22b-distilled.safetensors",
-    local_dir=str(weights_dir),
-    local_dir_use_symlinks=False,
-)
+# --- 1. 10Eros checkpoint: FP8→BF16 conversion + DEV metadata injection ---
+# Official DEV checkpoint is BF16. fp8_cast() expects BF16 input.
+# 10Eros is FP8 (CivitAI format) → convert to BF16 to match official dtype distribution.
+print("[1/4] Preparing 10Eros checkpoint...")
+eros_fp8_path = hf_hub_download(repo_id=EROS_REPO, filename=EROS_FILE)
+print(f"  Downloaded: {eros_fp8_path}")
+
+EROS_FIXED = "/tmp/eros_bf16_with_meta.safetensors"
+if os.path.exists(EROS_FIXED):
+    print("  Using cached BF16 checkpoint")
+else:
+    # Fetch DEV checkpoint metadata from header only (first 2MB, not full 46GB)
+    print("  Fetching DEV checkpoint metadata (header only)...")
+    dev_url = f"https://huggingface.co/{LTX_MODEL_REPO}/resolve/main/ltx-2.3-22b-dev.safetensors"
+    hdr_resp = requests.get(dev_url, headers={"Range": "bytes=0-2000000"}, timeout=30)
+    hdr_resp.raise_for_status()
+    hdr_size = int.from_bytes(hdr_resp.content[:8], "little")
+    hdr_json = json.loads(hdr_resp.content[8:8 + min(hdr_size, len(hdr_resp.content) - 8)])
+    dev_metadata = hdr_json.get("__metadata__", {})
+    print(f"  DEV metadata keys: {list(dev_metadata.keys())}")
+
+    # Convert FP8→BF16 + inject metadata
+    print("  Converting FP8→BF16 (lossless upcast)...")
+    _fp8_types = {torch.float8_e4m3fn, torch.float8_e5m2}
+    tensors = {}
+    _converted = 0
+    with safe_open(eros_fp8_path, framework="pt") as f:
+        for key in f.keys():
+            tensor = f.get_tensor(key)
+            if tensor.dtype in _fp8_types:
+                tensors[key] = tensor.to(torch.bfloat16)
+                _converted += 1
+            else:
+                tensors[key] = tensor
+    print(f"  Converted {_converted} FP8→BF16, kept {len(tensors)-_converted} as-is")
+    save_file(tensors, EROS_FIXED, metadata=dev_metadata)
+    del tensors
+    gc.collect()
+    print("  Saved with DEV metadata")
+
+checkpoint_path = EROS_FIXED
 spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.0.safetensors")
 gemma_root = snapshot_download(repo_id=GEMMA_REPO)
+distilled_lora_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=DISTILLED_LORA_FILE)
 
 # ---- Insert block (LoRA downloads) between lines 268 and 269 ----
 # LoRA repo + download the requested LoRA adapters
@@ -322,9 +362,16 @@ print(f"Demopose LoRA: {demopose_lora_path}")
 print(f"Checkpoint: {checkpoint_path}")
 print(f"Spatial upsampler: {spatial_upsampler_path}")
 print(f"Gemma root: {gemma_root}")
+print(f"Distilled LoRA: {distilled_lora_path}")
 
 # Initialize pipeline WITH text encoder and optional audio support
 # ---- Replace block (pipeline init) lines 275-281 ----
+print("Creating TI2VidTwoStagesPipeline...")
+distilled_lora = [LoraPathStrengthAndSDOps(
+    path=distilled_lora_path,
+    strength=0.9,
+    sd_ops=SDOps(name="distilled_lora", mapping=()),
+)]
 pipeline = LTX23DistilledA2VPipeline(
     distilled_checkpoint_path=checkpoint_path,
     spatial_upsampler_path=spatial_upsampler_path,