Update final_lipsync.py

final_lipsync.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import os
 import gc
+import urllib.request
 import cv2
 import torch
 import torch.nn as nn
@@ -25,18 +26,18 @@ from huggingface_hub import hf_hub_download
 import torch.serialization
 try:
     torch.serialization.add_safe_globals([np.core.multiarray.scalar])
-except:
+except Exception:
     pass
 
+# Robust torch.load patch — handles both positional and keyword args
 _original_load = torch.load
 def _patched_load(*args, **kwargs):
-    kwargs['weights_only'] = False
+    kwargs['weights_only'] = False  # Always force weights_only=False
     return _original_load(*args, **kwargs)
 torch.load = _patched_load
 # ----------------------------------------------
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# Secrets se API Key uthane ke liye
 OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
 
 # ==========================================
@@ -152,8 +153,48 @@ class MyFakeImageModel(nn.Module):
         self.fc = nn.Linear(512, 2)
     def forward(self, x): return torch.softmax(self.fc(x.view(x.size(0), -1)), dim=1)
 
+
+# ==========================================
+# 3. CLAP CHECKPOINT — DOWNLOAD HELPER
+# ==========================================
+def load_clap_safely(clap_module, device):
+    """
+    Load CLAP checkpoint with PyTorch 2.6+ compatibility.
+    Downloads checkpoint locally and uses strict=False to handle key mismatches.
+    """
+    CLAP_CKPT_URL = "https://huggingface.co/lukewys/laion_clap/resolve/main/music_audioset_epoch_15_esc_90.14.pt"
+    CLAP_CKPT_PATH = "/tmp/clap_music_audioset.pt"
+
+    # Step 1: Download checkpoint locally if not cached
+    if not os.path.exists(CLAP_CKPT_PATH):
+        print("⬇️ Downloading CLAP checkpoint to local disk...")
+        urllib.request.urlretrieve(CLAP_CKPT_URL, CLAP_CKPT_PATH)
+        print("✅ CLAP checkpoint downloaded!")
+    else:
+        print("✅ CLAP checkpoint already cached.")
+
+    # Step 2: Load checkpoint with weights_only=False
+    print("🔄 Loading CLAP checkpoint (strict=False)...")
+    ckpt = torch.load(CLAP_CKPT_PATH, map_location=device)
+
+    # Step 3: Extract state_dict from checkpoint
+    if isinstance(ckpt, dict):
+        if "state_dict" in ckpt:
+            state_dict = ckpt["state_dict"]
+        elif "model" in ckpt:
+            state_dict = ckpt["model"]
+        else:
+            state_dict = ckpt
+    else:
+        state_dict = ckpt
+
+    # Step 4: Load with strict=False to ignore unexpected/missing keys
+    clap_module.model.load_state_dict(state_dict, strict=False)
+    print("✅ CLAP model loaded successfully in lipsync!")
+
+
 # ==========================================
-# 3. LOAD ALL MODELS (REDIRECTION TO HUB)
+# 4. LOAD ALL MODELS (REDIRECTION TO HUB)
 # ==========================================
 def load_all_models():
     print("Loading Models...")
@@ -165,35 +206,45 @@ def load_all_models():
     m["ast_ext"] = AutoFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
     m["ast_mod"] = AutoModelForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593").to(DEVICE).eval()
 
+    # CLAP loading with PyTorch 2.6+ fix
     m["clap"] = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-tiny').to(DEVICE)
-    m["clap"].load_ckpt(model_id=1)
+    load_clap_safely(m["clap"], DEVICE)
 
     sync_m = SyncNet_color().to(DEVICE)
-    try: sync_m.load_state_dict(torch.load(hf_hub_download(repo_id="camenduru/Wav2Lip", filename="syncnet_v2.pth"), map_location=DEVICE)['state_dict'], strict=False)
-    except: pass
+    try:
+        sync_m.load_state_dict(torch.load(hf_hub_download(repo_id="camenduru/Wav2Lip", filename="syncnet_v2.pth"), map_location=DEVICE)['state_dict'], strict=False)
+    except Exception as e:
+        print(f"⚠️ SyncNet load issue: {e}")
     m["hf_sync"] = sync_m.eval()
 
     # Step 8 Models - Downloading from your HuggingFace Repo
     img_m = MyFakeImageModel().to(DEVICE)
     try:
-        # Aapka FAKE-IMAGE model repo link use karein
         img_path = hf_hub_download(repo_id="aneela-pervez/My-Deepfake-Models", filename="checkpoint_step000080000.pth")
         img_m.load_state_dict(torch.load(img_path, map_location=DEVICE), strict=False)
-    except: pass
+    except Exception as e:
+        print(f"⚠️ Custom image model load issue: {e}")
     m["custom_img"] = img_m.eval()
 
     fusion_m = FusionLipSyncModel().to(DEVICE)
     try:
         fusion_path = hf_hub_download(repo_id="aneela-pervez/My-Deepfake-Models", filename="best_model.pth")
-        fusion_m.load_state_dict(torch.load(fusion_path, map_location=DEVICE).get('model_state_dict'), strict=False)
-    except: pass
+        fusion_state = torch.load(fusion_path, map_location=DEVICE)
+        if isinstance(fusion_state, dict) and 'model_state_dict' in fusion_state:
+            fusion_state = fusion_state['model_state_dict']
+        fusion_m.load_state_dict(fusion_state, strict=False)
+    except Exception as e:
+        print(f"⚠️ Fusion model load issue: {e}")
     m["custom_fusion"] = fusion_m.eval()
 
     return m
 
 models = load_all_models()
 
-# (Analysis logic remains same as per your snippet)
+
+# ==========================================
+# 5. ANALYSIS LOGIC (UNCHANGED)
+# ==========================================
 def extract_mouth(frame_np):
     try: faces = RetinaFace.detect_faces(frame_np)
     except: return None
@@ -207,18 +258,18 @@ def full_analysis(video_path):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
     frames = []
-    
+
     # Fast Processing Logic (Har 5th frame use karega taake GPU pe time bache)
     frame_count = 0
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret: break
-        
+
         if frame_count % 5 == 0:
             temp_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             small_frame = cv2.resize(temp_frame, (640, 360))
             frames.append(small_frame)
-            
+
         frame_count += 1
         if len(frames) > 150: break
     cap.release()
@@ -240,7 +291,7 @@ def full_analysis(video_path):
     # Fast Audio Loading Logic
     y, sr = librosa.load(video_path, sr=48000, duration=10)
     y_16, _ = librosa.load(video_path, sr=16000, duration=10)
-    
+
     with torch.no_grad():
         a_emb = torch.from_numpy(models["clap"].get_audio_embedding_from_data(x=[y])).to(DEVICE)
         r_a_emb = torch.from_numpy(models["clap"].get_text_embedding(PROMPTS["REAL_AUDIO"])).to(DEVICE)
@@ -279,9 +330,9 @@ def master_pipeline(video_path):
     res = full_analysis(video_path)
     if res is None or res[0] is None:
         return None, "Analysis failed to process video.", "Error"
-        
+
     img, v_score, a_score, sync_score, clip_v, cust_v, clap_a, ast_a, hf_s, cust_s = res
-    
+
     is_v_fake, is_a_fake, is_sync_bad = v_score >= 50, a_score >= 50, sync_score >= 60
     if is_v_fake and is_a_fake: case = "CASE 1: Full Deep Fake"
     elif is_v_fake: case = "CASE 2: Fake Video + Real Audio"