Update app.py

app.py

@@ -8,8 +8,9 @@ import os
 
 import torch.nn.functional as F
 
-# Check for CUDA availability but fallback to CPU
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Force CPU mode for Zero GPU environment
+device = torch.device('cpu')
+torch.set_num_threads(4)  # Optimize CPU performance
 
 norm_layer = nn.InstanceNorm2d
 
@@ -86,28 +87,41 @@ class Generator(nn.Module):
         out = self.model4(out)
         return out
 
-# Initialize models
+# Initialize models with error handling and memory optimization
 def load_models():
     try:
-        model1 = Generator(3, 1, 3).to(device)
-        model2 = Generator(3, 1, 3).to(device)
+        print("Initializing models in CPU mode...")
         
-        # 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))
+        # Initialize models
+        model1 = Generator(3, 1, 3)
+        model2 = Generator(3, 1, 3)
         
+        # Load model weights with explicit CPU mapping
+        model1.load_state_dict(torch.load('model.pth', map_location='cpu'))
+        model2.load_state_dict(torch.load('model2.pth', map_location='cpu'))
+        
+        # Set to eval mode and optimize for inference
         model1.eval()
         model2.eval()
+        
+        # Enable inference optimizations
+        torch.set_grad_enabled(False)
+        
+        print("Models loaded successfully in CPU mode")
         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.")
+        error_msg = f"Error loading models: {str(e)}"
+        print(error_msg)
+        raise gr.Error("Failed to initialize models. Please check the model files and system configuration.")
 
+# Load models with proper error handling
 try:
+    print("Starting model initialization...")
     model1, model2 = load_models()
+    print("Model initialization completed")
 except Exception as e:
-    print(f"Model initialization failed: {str(e)}")
-    model1 = model2 = None
+    print(f"Critical error during model initialization: {str(e)}")
+    raise gr.Error("Failed to start the application due to model initialization error.")
 
 def apply_style_transfer(img, strength=1.0):
     """Apply artistic style transfer effect"""
@@ -287,4 +301,10 @@ with gr.Blocks(css=custom_css) as iface:
     )
 
 # Launch the interface
-iface.launch()
+iface.launch(
+    server_name="0.0.0.0",  # Required for Hugging Face Spaces
+    server_port=7860,       # Default port for Hugging Face Spaces
+    share=False,            # Disable public URL
+    debug=False,            # Disable debug mode
+    enable_queue=True       # Enable queue for better performance
+)