Update app.py

app.py

@@ -7,108 +7,82 @@ import os
 
 print("🚀 Starting ZeroIG app...")
 
-class ZeroIGEnhancer:
-    def __init__(self):
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        print(f"Using device: {self.device}")
-        self.model = self.load_model()
+# Load ZeroIG model
+def load_zeroig():
+    try:
+        from model import Finetunemodel, Network
+        
+        model_path = "./weights/model.pt"
+        if os.path.exists(model_path):
+            model = Finetunemodel(model_path)
+            print("✅ Loaded trained ZeroIG model")
+        else:
+            model = Network()
+            print("⚠️ Using ZeroIG with random weights")
+        
+        model.eval()
+        return model
+    except Exception as e:
+        print(f"⚠️ Could not load ZeroIG: {e}")
+        return None
+
+# Initialize model
+zeroig_model = load_zeroig()
+
+def enhance_image(image):
+    """Enhance low-light image"""
+    if image is None:
+        return None
     
-    def load_model(self):
-        try:
-            from model import Finetunemodel, Network
+    try:
+        if zeroig_model is not None:
+            # Resize if too large
+            if max(image.size) > 800:
+                ratio = 800 / max(image.size)
+                new_size = tuple(int(dim * ratio) for dim in image.size)
+                image = image.resize(new_size, Image.Resampling.LANCZOS)
             
-            # Check for trained weights
-            model_path = "./weights/model.pt"
-            if os.path.exists(model_path):
-                model = Finetunemodel(model_path)
-                print("✅ Loaded trained ZeroIG model")
-            else:
-                model = Network()
-                print("⚠️ Using ZeroIG with random weights")
+            # Convert to tensor
+            transform = transforms.ToTensor()
+            input_tensor = transform(image).unsqueeze(0)
             
-            model.to(self.device)
-            model.eval()
-            return model
+            with torch.no_grad():
+                if hasattr(zeroig_model, 'enhance'):
+                    # Finetunemodel
+                    enhanced, denoised = zeroig_model(input_tensor)
+                    result = denoised
+                else:
+                    # Network model
+                    outputs = zeroig_model(input_tensor)
+                    result = outputs[13]  # H3 output
             
-        except Exception as e:
-            print(f"⚠️ Could not load ZeroIG model: {e}")
-            return None
-    
-    def enhance_image(self, image):
-        if image is None:
-            return None
+            # Convert back to image
+            result = result.squeeze(0).clamp(0, 1)
+            to_pil = transforms.ToPILImage()
+            return to_pil(result)
         
-        try:
-            if self.model is not None:
-                # Resize image if too large to prevent memory issues
-                max_size = 800
-                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)
-                
-                # Convert to tensor
-                transform = transforms.ToTensor()
-                input_tensor = transform(image).unsqueeze(0).to(self.device)
-                
-                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
-                    else:
-                        # Network model - returns multiple outputs
-                        outputs = self.model(input_tensor)
-                        result_tensor = outputs[13]  # H3 is the final denoised result
-                
-                # Convert back to PIL
-                result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1)
-                to_pil = transforms.ToPILImage()
-                return to_pil(result_tensor)
-            
-            else:
-                # Fallback: simple enhancement
-                img_array = np.array(image).astype(np.float32)
-                enhanced = np.clip(img_array * 1.8, 0, 255).astype(np.uint8)
-                return Image.fromarray(enhanced)
-                
-        except Exception as e:
-            print(f"Enhancement error: {e}")
-            # Fallback enhancement
-            img_array = np.array(image).astype(np.float32)
-            enhanced = np.clip(img_array * 1.5, 0, 255).astype(np.uint8)
+        else:
+            # Simple fallback
+            img_array = np.array(image)
+            enhanced = np.clip(img_array * 1.8, 0, 255).astype(np.uint8)
             return Image.fromarray(enhanced)
+            
+    except Exception as e:
+        print(f"Error: {e}")
+        # Emergency fallback
+        img_array = np.array(image)
+        enhanced = np.clip(img_array * 1.5, 0, 255).astype(np.uint8)
+        return Image.fromarray(enhanced)
 
-# Initialize the enhancer
-enhancer = ZeroIGEnhancer()
-
-def process_image(image):
-    """Process uploaded image"""
-    return enhancer.enhance_image(image)
-
-# Create Gradio interface - simplified for HF Spaces
-print("Creating Gradio interface...")
+# Create interface with minimal parameters
+print("Creating interface...")
 
-interface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil", label="📷 Upload Low-Light Image"),
-    outputs=gr.Image(type="pil", label="✨ Enhanced Image"),
-    title="🌟 ZeroIG: Zero-Shot Low-Light Enhancement",
-    description="""
-    **CVPR 2024** - Zero-shot illumination-guided joint denoising and adaptive enhancement for low-light images.
-    
-    🔸 Upload a low-light or dark image  
-    🔸 Get enhanced version with noise reduction  
-    🔸 No training data required!  
-    
-    📄 [Paper](https://openaccess.thecvf.com/content/CVPR2024/papers/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_Joint_Denoising_and_Adaptive_Enhancement_for_Low-Light_CVPR_2024_paper.pdf) | 💻 [Code](https://github.com/Doyle59217/ZeroIG)
-    """,
-    examples=None,
-    cache_examples=False,
-    flagging_mode="never"
+demo = gr.Interface(
+    fn=enhance_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="ZeroIG: Low-Light Enhancement"
 )
 
-# Launch - configured specifically for Hugging Face Spaces
-if __name__ == "__main__":
-    print("🚀 Launching Gradio interface...")
-    interface.launch()
+print("🚀 Launching...")
+demo.launch()