Update app.py

app.py

@@ -8,14 +8,47 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
+import os
+import urllib.request
+import zipfile
+
 # --- CONFIG & MODEL DOWNLOAD ---
 MODEL_PATH = "LookThem_V8_MNIST.pth"
+ZIP_PATH = "LookThem_V8_MNIST.zip"
 HF_URL = "https://huggingface.co/ASomeoneWhoInterestedWithAI/LookThem_V8-MNIST_Classifier/resolve/main/LookThem_V8_MNIST%20(2).pth"
 
 if not os.path.exists(MODEL_PATH):
     print(f"Downloading model weights from Hugging Face...")
-    urllib.request.urlretrieve(HF_URL, MODEL_PATH)
-    print("Download complete!")
+    # Download the file as a zip first
+    urllib.request.urlretrieve(HF_URL, ZIP_PATH)
+    print("Download complete! Checking for zip compression...")
+    
+    try:
+        # Unzip the file
+        with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
+            # Look for a .pth file inside the zip
+            file_list = zip_ref.namelist()
+            pth_files = [f for f in file_list if f.endswith('.pth')]
+            
+            if pth_files:
+                # Extract the .pth file and rename it to our expected MODEL_PATH
+                zip_ref.extract(pth_files[0], path=".")
+                if pth_files[0] != MODEL_PATH:
+                    os.rename(pth_files[0], MODEL_PATH)
+                print(f"Successfully extracted: {pth_files[0]} -> {MODEL_PATH}")
+            else:
+                # If no .pth inside, maybe the zip itself *is* the model (PyTorch 2.0+ format)
+                print("No .pth file found inside zip. Renaming zip directly to .pth...")
+                os.rename(ZIP_PATH, MODEL_PATH)
+                
+    except zipfile.BadZipFile:
+        # If it wasn't actually a zip file, just rename it
+        print("File is not a zip archive. Proceeding with standard weight loading.")
+        os.rename(ZIP_PATH, MODEL_PATH)
+        
+    # Clean up the temporary zip file if it still exists
+    if os.path.exists(ZIP_PATH):
+        os.remove(ZIP_PATH)
 
 # --- DEFINE YOUR MODEL ARCHITECTURE ---
 class LookThemLayer(nn.Module):