Update app.py

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 from PIL import Image
 import torchvision.transforms as transforms
 import os
-from huggingface_hub import hf_hub_download
+
 import torch.nn.functional as F
 
 # Check for CUDA availability but fallback to CPU
@@ -88,21 +88,26 @@ class Generator(nn.Module):
 
 # Initialize models
 def load_models():
-    model1 = Generator(3, 1, 3).to(device)
-    model2 = Generator(3, 1, 3).to(device)
-    
-    # Download models from HuggingFace Hub
-    model1_path = hf_hub_download(repo_id="your-hf-repo/line-drawing", filename="model.pth")
-    model2_path = hf_hub_download(repo_id="your-hf-repo/line-drawing", filename="model2.pth")
-    
-    model1.load_state_dict(torch.load(model1_path, map_location=device))
-    model2.load_state_dict(torch.load(model2_path, map_location=device))
-    
-    model1.eval()
-    model2.eval()
-    return model1, model2
+    try:
+        model1 = Generator(3, 1, 3).to(device)
+        model2 = Generator(3, 1, 3).to(device)
+        
+        # Load local model files
+        model1.load_state_dict(torch.load('model.pth', map_location=device))
+        model2.load_state_dict(torch.load('model2.pth', map_location=device))
+        
+        model1.eval()
+        model2.eval()
+        return model1, model2
+    except Exception as e:
+        print(f"Error loading models: {str(e)}")
+        raise gr.Error("Failed to load models. Please check if model files exist in the correct location.")
 
-model1, model2 = load_models()
+try:
+    model1, model2 = load_models()
+except Exception as e:
+    print(f"Model initialization failed: {str(e)}")
+    model1 = model2 = None
 
 def apply_style_transfer(img, strength=1.0):
     """Apply artistic style transfer effect"""