added files for inference

.gitignore

@@ -0,0 +1,37 @@
+# Large model files (stored on Hugging Face)
+checkpoints/
+pretrained_weights/
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+dist/
+build/
+
+# Environment
+.env
+.venv
+env/
+venv/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Logs
+*.log
+

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 import torchvision.transforms.functional as TF
 from pathlib import Path
 import sys
+from huggingface_hub import hf_hub_download
 
 # Add src to path
 sys.path.insert(0, str(Path(__file__).parent / "src"))
@@ -17,6 +18,9 @@ from autoencoder import get_vqgan
 from noiseControl import resshift_schedule
 from config import device, T, k, normalize_input, latent_flag, gt_size
 
+# Hugging Face repo ID for weights
+HF_WEIGHTS_REPO_ID = "shekkari21/DiffusionSR-weights"
+
 # Global variables for loaded models
 model = None
 autoencoder = None
@@ -30,9 +34,12 @@ def load_models():
     print("Loading models...")
     
     # Load model checkpoint
-    checkpoint_path = "checkpoints/ckpts/model_3200.pth"  # Update with your checkpoint path
-    if not Path(checkpoint_path).exists():
-        # Try to find any checkpoint
+    checkpoint_path = "checkpoints/ckpts/model_3200.pth"
+    checkpoint_file = Path(checkpoint_path)
+    
+    # Download from Hugging Face if not found locally
+    if not checkpoint_file.exists():
+        # Try to find any checkpoint locally first
         ckpt_dir = Path("checkpoints/ckpts")
         if ckpt_dir.exists():
             checkpoints = list(ckpt_dir.glob("model_*.pth"))
@@ -40,9 +47,42 @@ def load_models():
                 checkpoint_path = str(checkpoints[-1])  # Use latest
                 print(f"Using checkpoint: {checkpoint_path}")
             else:
-                raise FileNotFoundError("No model checkpoint found!")
+                # Download from Hugging Face
+                print(f"Checkpoint not found locally. Downloading from Hugging Face...")
+                try:
+                    # Files are in root of weights repo, download to local directory structure
+                    ckpt_dir.mkdir(parents=True, exist_ok=True)
+                    downloaded_path = hf_hub_download(
+                        repo_id=HF_WEIGHTS_REPO_ID,
+                        filename="model_3200.pth",
+                        local_dir=str(ckpt_dir),
+                        local_dir_use_symlinks=False
+                    )
+                    checkpoint_path = str(ckpt_dir / "model_3200.pth")
+                    print(f"✓ Downloaded checkpoint: {checkpoint_path}")
+                except Exception as e:
+                    raise FileNotFoundError(
+                        f"Could not download checkpoint from Hugging Face: {e}\n"
+                        f"Please ensure the file exists in the repository."
+                    )
         else:
-            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+            # Create directory and download
+            ckpt_dir.mkdir(parents=True, exist_ok=True)
+            print(f"Checkpoint not found locally. Downloading from Hugging Face...")
+            try:
+                downloaded_path = hf_hub_download(
+                    repo_id=HF_WEIGHTS_REPO_ID,
+                    filename="model_3200.pth",
+                    local_dir=str(ckpt_dir),
+                    local_dir_use_symlinks=False
+                )
+                checkpoint_path = str(ckpt_dir / "model_3200.pth")
+                print(f"✓ Downloaded checkpoint: {checkpoint_path}")
+            except Exception as e:
+                raise FileNotFoundError(
+                    f"Could not download checkpoint from Hugging Face: {e}\n"
+                    f"Please ensure the file exists in the repository."
+                )
     
     model = FullUNET()
     model = model.to(device)

requirements.txt

@@ -10,4 +10,5 @@ lpips>=0.1.4
 loralib>=0.1.2
 python-dotenv>=1.0.0
 numpy>=1.24.0
+huggingface_hub>=0.20.0
 

src/autoencoder.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 from pathlib import Path
 import sys
 import os
+from huggingface_hub import hf_hub_download
 
 # Handle import of ldm from latent-diffusion repository
 # Check if ldm directory exists locally (from latent-diffusion repo)
@@ -49,6 +50,9 @@ from config import (
     device
 )
 
+# Hugging Face repo ID for weights
+HF_WEIGHTS_REPO_ID = "shekkari21/DiffusionSR-weights"
+
 
 def load_vqgan(ckpt_path=None, device=device):
     """
@@ -68,8 +72,29 @@ def load_vqgan(ckpt_path=None, device=device):
     if not Path(ckpt_path).is_absolute():
         ckpt_path = _project_root / ckpt_path
     
+    # Download from Hugging Face if not found locally
     if not Path(ckpt_path).exists():
-        raise FileNotFoundError(f"VQGAN checkpoint not found at: {ckpt_path}")
+        print(f"VQGAN checkpoint not found locally. Downloading from Hugging Face...")
+        try:
+            # Files are in root of weights repo, download to local directory structure
+            local_weights_dir = _project_root / "pretrained_weights"
+            local_weights_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Download from root of weights repo
+            downloaded_path = hf_hub_download(
+                repo_id=HF_WEIGHTS_REPO_ID,
+                filename="autoencoder_vq_f4.pth",
+                local_dir=str(local_weights_dir),
+                local_dir_use_symlinks=False
+            )
+            ckpt_path = local_weights_dir / "autoencoder_vq_f4.pth"
+            print(f"✓ Downloaded VQGAN checkpoint: {ckpt_path}")
+        except Exception as e:
+            raise FileNotFoundError(
+                f"VQGAN checkpoint not found at: {ckpt_path}\n"
+                f"Could not download from Hugging Face: {e}\n"
+                f"Please ensure the file exists in the repository."
+            )
     
     print(f"Loading VQGAN from: {ckpt_path}")