Update app.py

app.py

@@ -6,6 +6,7 @@ import base64
 import torch
 import torchvision.transforms as transforms
 import os
+import sys
 
 app = Flask(__name__)
 
@@ -90,68 +91,84 @@ class ZeroIGProcessor:
     
     def load_model(self):
         try:
+            # Import your uploaded ZeroIG files
             from model import Finetunemodel, Network
             
             # Try to load trained weights
             model_path = "./weights/model.pt"
             if os.path.exists(model_path):
+                print(f"Found model weights at {model_path}")
                 model = Finetunemodel(model_path)
-                print("✅ Loaded trained ZeroIG model")
+                print("✅ Loaded ZeroIG Finetunemodel with trained weights")
             else:
+                print("No trained weights found, using Network with random initialization")
                 model = Network()
-                print("⚠️ Using ZeroIG with random weights (train a model for better results)")
+                print("⚠️ Using ZeroIG Network with random weights")
             
             model.to(self.device)
             model.eval()
+            print(f"Model moved to {self.device}")
             return model
             
+        except ImportError as e:
+            print(f"❌ Could not import ZeroIG modules: {e}")
+            print("Make sure you have uploaded: model.py, loss.py, utils.py")
+            return None
         except Exception as e:
             print(f"❌ Could not load ZeroIG model: {e}")
             return None
     
     def enhance_image(self, image):
-        """Enhance image using ZeroIG model"""
+        """Enhance image using your ZeroIG model"""
         try:
             if self.model is None:
-                return self.simple_enhance(image), "Using simple enhancement (ZeroIG model not available)"
+                return self.simple_enhance(image), "❌ ZeroIG model not available - using simple enhancement"
             
             # Resize if too large to prevent memory issues
             original_size = image.size
-            max_size = 1024
+            max_size = 800  # Adjust based on your needs
             if max(image.size) > max_size:
                 ratio = max_size / max(image.size)
                 new_size = tuple(int(dim * ratio) for dim in image.size)
                 image = image.resize(new_size, Image.Resampling.LANCZOS)
+                print(f"Resized image from {original_size} to {image.size}")
             
             # Convert to tensor
             transform = transforms.ToTensor()
             input_tensor = transform(image).unsqueeze(0).to(self.device)
+            print(f"Input tensor shape: {input_tensor.shape}")
             
-            # Run ZeroIG model
+            # Run your ZeroIG model
             with torch.no_grad():
                 if hasattr(self.model, 'enhance') and hasattr(self.model, 'denoise_1'):
                     # Finetunemodel - returns (enhanced, denoised)
                     enhanced, denoised = self.model(input_tensor)
-                    result_tensor = denoised
-                    status = "✅ Enhanced with ZeroIG Finetuned model"
+                    result_tensor = denoised  # Use denoised output
+                    status = "✅ Enhanced with ZeroIG Finetunemodel"
+                    print("Used Finetunemodel")
                 else:
                     # Network model - returns multiple outputs
                     outputs = self.model(input_tensor)
                     result_tensor = outputs[13]  # H3 is the final denoised result
                     status = "✅ Enhanced with ZeroIG Network model"
+                    print("Used Network model")
             
             # Convert back to PIL
             result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1)
             enhanced_image = transforms.ToPILImage()(result_tensor)
+            print(f"Output image size: {enhanced_image.size}")
             
             # Resize back to original size if needed
-            if enhanced_image.size != original_size:
+            if enhanced_image.size != original_size and original_size != image.size:
                 enhanced_image = enhanced_image.resize(original_size, Image.Resampling.LANCZOS)
+                print(f"Resized back to original size: {enhanced_image.size}")
             
             return enhanced_image, status
             
         except Exception as e:
             print(f"ZeroIG enhancement error: {e}")
+            import traceback
+            traceback.print_exc()
             return self.simple_enhance(image), f"⚠️ ZeroIG failed, using simple enhancement: {str(e)}"
     
     def simple_enhance(self, image):
@@ -191,7 +208,7 @@ def index():
                 # Store original for comparison
                 original_image = image_to_base64(image)
                 
-                # Enhance with ZeroIG
+                # Enhance with your ZeroIG model
                 enhanced_image, enhancement_status = zeroig.enhance_image(image)
                 
                 # Convert result to base64
@@ -201,6 +218,8 @@ def index():
         except Exception as e:
             error = f"Error processing image: {str(e)}"
             print(f"Error: {e}")
+            import traceback
+            traceback.print_exc()
     
     return render_template_string(HTML_TEMPLATE, 
                                   original_image=original_image,