Upload 2 files

app.py

@@ -81,16 +81,24 @@ def ensure_data_ready(data_dir: Path, checkpoints: list) -> bool:
     """Ensure data and checkpoints are downloaded."""
     print(f"Checking for existing data in {data_dir.absolute()}...")
 
-    # Check which models have data (config + embeddings)
+    # Check which models have data (config + embeddings + images)
     models_with_data = []
     for model in ["dmd2", "edm"]:
         config_path = data_dir / model / "config.json"
         embeddings_dir = data_dir / model / "embeddings"
-        if config_path.exists() and embeddings_dir.exists():
-            csv_files = list(embeddings_dir.glob("*.csv"))
-            if csv_files:
-                models_with_data.append(model)
-                print(f"  Found {model}: {config_path}")
+        images_dir = data_dir / model / "images" / "imagenet_real"
+
+        if not config_path.exists():
+            continue
+        if not embeddings_dir.exists():
+            continue
+
+        csv_files = list(embeddings_dir.glob("*.csv"))
+        png_files = list(images_dir.glob("sample_*.png")) if images_dir.exists() else []
+
+        if csv_files and png_files:
+            models_with_data.append(model)
+            print(f"  Found {model}: {len(csv_files)} csv, {len(png_files)} images")
 
     if not models_with_data:
         print("Data not found, downloading...")

requirements.txt

@@ -5,20 +5,22 @@ git+https://github.com/mckellcarter/diffviews.git
 # Core dependencies
 torch>=2.0.0
 numpy>=1.21.0
-numba==0.58.1 
 pandas>=1.5.0
 pillow>=9.0.0
 scikit-learn>=1.0.0
 umap-learn>=0.5.0
 tqdm>=4.60.0
 
+# Numba pinned for UMAP pickle compatibility
+numba==0.58.1
+
 # Visualization
 gradio>=4.0.0
 plotly>=5.18.0
 matplotlib>=3.5.0
 
-# HuggingFace Hub for data download
-huggingface_hub>=0.19.0,<0.28.0 
+# HuggingFace Hub - pinned for Gradio 6 compatibility
+huggingface_hub>=0.19.0,<0.28.0
 
 # Optional but useful
 scipy>=1.7.0