Upload app.py with huggingface_hub

app.py

@@ -1,7 +1,11 @@
-"""E3Diff: SAR-to-Optical Translation - HuggingFace Space."""
+"""
+E3Diff: SAR-to-Optical Translation - HuggingFace Space
+Exact copy of working local implementation
+"""
 
 import os
 import torch
+import torch.nn as nn
 import numpy as np
 from PIL import Image, ImageEnhance
 import gradio as gr
@@ -9,7 +13,7 @@ import tempfile
 import time
 from huggingface_hub import hf_hub_download
 
-# Import model components
+# Import model components (exact same as local)
 from unet import UNet
 from diffusion import GaussianDiffusion
 
@@ -23,36 +27,47 @@ except ImportError:
 
 
 class E3DiffInference:
-    """E3Diff Inference Pipeline - matches local implementation exactly."""
+    """
+    E3Diff Inference Pipeline - EXACT copy from local inference.py
+    """
     
-    def __init__(self, weights_path=None, device="cuda", num_inference_steps=1):
+    def __init__(self, weights_path=None, device="cuda", image_size=256, num_inference_steps=1):
         self.device = torch.device(device if torch.cuda.is_available() else "cpu")
-        self.image_size = 256
+        self.image_size = image_size
         self.num_inference_steps = num_inference_steps
         
         print(f"[E3Diff] Initializing on device: {self.device}")
+        print(f"[E3Diff] Image size: {image_size}x{image_size}")
         print(f"[E3Diff] Inference steps: {num_inference_steps}")
         
+        # Build model
         self.model = self._build_model()
+        
+        # Load weights
         self._load_weights(weights_path)
+        
+        # Set to eval mode
         self.model.eval()
-        print("[E3Diff] Model ready!")
+        
+        print("[E3Diff] Model ready for inference!")
     
     def _build_model(self):
-        """Build model - exact same config as local inference.py"""
+        """Build the E3Diff model architecture - exact same config."""
+        # UNet configuration from SEN12_256_s2_test.json
         unet = UNet(
-            in_channel=3,
-            out_channel=3,
+            in_channel=3,           # Noisy image channels
+            out_channel=3,          # Output optical image
             norm_groups=16,
             inner_channel=64,
-            channel_mults=[1, 2, 4, 8, 16],
-            attn_res=[],
+            channel_mults=[1, 2, 4, 8, 16],  # Encoder/decoder channels
+            attn_res=[],            # No attention at specific resolutions
             res_blocks=1,
             dropout=0,
             image_size=self.image_size,
-            condition_ch=3
+            condition_ch=3          # SAR condition channels
         )
         
+        # Diffusion wrapper
         schedule_opt = {
             'schedule': 'linear',
             'n_timestep': self.num_inference_steps,
@@ -88,7 +103,7 @@ class E3DiffInference:
         return model.to(self.device)
     
     def _load_weights(self, weights_path):
-        """Load weights - same as local inference.py"""
+        """Load pre-trained weights."""
         if weights_path is None:
             weights_path = hf_hub_download(
                 repo_id="Dhenenjay/E3Diff-SAR2Optical",
@@ -98,38 +113,48 @@ class E3DiffInference:
         print(f"[E3Diff] Loading weights from: {weights_path}")
         state_dict = torch.load(weights_path, map_location=self.device, weights_only=False)
         self.model.load_state_dict(state_dict, strict=False)
-        print("[E3Diff] Weights loaded!")
+        print(f"[E3Diff] Weights loaded successfully!")
     
     def preprocess(self, image):
-        """Preprocess input image."""
+        """Preprocess input SAR image."""
+        # Convert to RGB if grayscale
         if image.mode != 'RGB':
             image = image.convert('RGB')
+        
+        # Resize to model input size
         if image.size != (self.image_size, self.image_size):
             image = image.resize((self.image_size, self.image_size), Image.LANCZOS)
         
+        # Convert to tensor and normalize to [-1, 1]
         img_np = np.array(image).astype(np.float32) / 255.0
-        img_tensor = torch.from_numpy(img_np).permute(2, 0, 1)
-        img_tensor = img_tensor * 2.0 - 1.0
+        img_tensor = torch.from_numpy(img_np).permute(2, 0, 1)  # HWC -> CHW
+        img_tensor = img_tensor * 2.0 - 1.0  # [0,1] -> [-1,1]
+        
         return img_tensor.unsqueeze(0).to(self.device)
     
     def postprocess(self, tensor):
-        """Postprocess output tensor."""
+        """Postprocess output tensor to PIL Image."""
+        # Clamp and denormalize
         tensor = tensor.squeeze(0).cpu()
         tensor = torch.clamp(tensor, -1, 1)
-        tensor = (tensor + 1.0) / 2.0
+        tensor = (tensor + 1.0) / 2.0  # [-1,1] -> [0,1]
+        
+        # Convert to numpy and PIL
         img_np = (tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
         return Image.fromarray(img_np)
     
     @torch.no_grad()
     def translate(self, sar_image, seed=42):
-        """Translate SAR to optical - same as local inference.py"""
+        """Translate SAR image to optical image."""
+        # Set seed for reproducibility
         if seed is not None:
             torch.manual_seed(seed)
             np.random.seed(seed)
         
-        sar_tensor = self.preprocess(sar_image)
+        # Preprocess
+        sar_tensor = self.preprocess(sar_image)  # [1, 3, H, W]
         
-        # Set noise schedule
+        # Set noise schedule for inference
         self.model.set_new_noise_schedule(
             {
                 'schedule': 'linear',
@@ -144,25 +169,30 @@ class E3DiffInference:
         )
         
         # Run inference
-        output, _ = self.model.super_resolution(sar_tensor, continous=False, seed=seed, img_s1=sar_tensor)
+        output, output_onestep = self.model.super_resolution(
+            sar_tensor,
+            continous=False,
+            seed=seed if seed is not None else 1,
+            img_s1=sar_tensor
+        )
+        
         return self.postprocess(output)
 
 
 class HighResProcessor:
-    """High resolution tiled processing."""
+    """High resolution tiled processing - exact copy from process_highres.py"""
     
     def __init__(self, device="cuda"):
         self.device = device
         self.model = None
         self.tile_size = 256
-        self.num_steps = None
     
-    def load_model(self, num_steps=1):
-        print(f"Loading E3Diff model with {num_steps} steps...")
-        self.model = E3DiffInference(device=self.device, num_inference_steps=num_steps)
-        self.num_steps = num_steps
+    def load_model(self):
+        print("Loading E3Diff model...")
+        self.model = E3DiffInference(device=self.device, num_inference_steps=1)
     
     def create_blend_weights(self, tile_size, overlap):
+        """Create smooth blending weights for seamless output."""
         ramp = np.linspace(0, 1, overlap)
         weight = np.ones((tile_size, tile_size))
         weight[:overlap, :] *= ramp[:, np.newaxis]
@@ -171,9 +201,10 @@ class HighResProcessor:
         weight[:, -overlap:] *= ramp[np.newaxis, ::-1]
         return weight[:, :, np.newaxis]
     
-    def process(self, image, overlap=64, num_steps=1):
-        if self.model is None or self.num_steps != num_steps:
-            self.load_model(num_steps)
+    def process(self, image, overlap=64):
+        """Process image at full resolution with seamless tiling."""
+        if self.model is None:
+            self.load_model()
         
         if isinstance(image, Image.Image):
             if image.mode != 'RGB':
@@ -186,49 +217,59 @@ class HighResProcessor:
         tile_size = self.tile_size
         step = tile_size - overlap
         
+        # Pad image
         pad_h = (step - (h - overlap) % step) % step
         pad_w = (step - (w - overlap) % step) % step
         img_padded = np.pad(img_np, ((0, pad_h), (0, pad_w), (0, 0)), mode='reflect')
         
         h_pad, w_pad = img_padded.shape[:2]
         
+        # Output arrays
         output = np.zeros((h_pad, w_pad, 3), dtype=np.float32)
         weights = np.zeros((h_pad, w_pad, 1), dtype=np.float32)
         blend_weight = self.create_blend_weights(tile_size, overlap)
         
+        # Calculate positions
         y_positions = list(range(0, h_pad - tile_size + 1, step))
         x_positions = list(range(0, w_pad - tile_size + 1, step))
         total_tiles = len(y_positions) * len(x_positions)
         
-        print(f"Processing {total_tiles} tiles at {w}x{h}...")
+        print(f"Processing {total_tiles} tiles ({len(x_positions)}x{len(y_positions)}) at {w}x{h}...")
         
         tile_idx = 0
         for y in y_positions:
             for x in x_positions:
+                # Extract tile
                 tile = img_padded[y:y+tile_size, x:x+tile_size]
                 tile_pil = Image.fromarray((tile * 255).astype(np.uint8))
                 
+                # Translate
                 result_pil = self.model.translate(tile_pil, seed=42)
                 result = np.array(result_pil).astype(np.float32) / 255.0
                 
+                # Blend
                 output[y:y+tile_size, x:x+tile_size] += result * blend_weight
                 weights[y:y+tile_size, x:x+tile_size] += blend_weight
                 
                 tile_idx += 1
-                if tile_idx % 4 == 0 or tile_idx == total_tiles:
+                if tile_idx % 10 == 0 or tile_idx == total_tiles:
                     print(f"  Tile {tile_idx}/{total_tiles}")
         
+        # Normalize
         output = output / (weights + 1e-8)
         output = output[:h, :w]
         
         return (output * 255).astype(np.uint8)
     
-    def enhance(self, image, contrast=1.1, sharpness=1.15, color=1.1):
+    def enhance(self, image, contrast=1.1, sharpness=1.2, color=1.1):
+        """Professional post-processing."""
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
+        
         image = ImageEnhance.Contrast(image).enhance(contrast)
         image = ImageEnhance.Sharpness(image).enhance(sharpness)
         image = ImageEnhance.Color(image).enhance(color)
+        
         return image
 
 
@@ -259,7 +300,7 @@ def load_sar_image(filepath):
     return Image.open(filepath).convert('RGB')
 
 
-def _translate_sar_impl(file, num_steps, overlap, enhance_output):
+def _translate_sar_impl(file, overlap, enhance_output):
     """Main translation function."""
     global processor
     
@@ -278,7 +319,7 @@ def _translate_sar_impl(file, num_steps, overlap, enhance_output):
     print(f"Input size: {w}x{h}")
     
     start = time.time()
-    result = processor.process(image, overlap=int(overlap), num_steps=int(num_steps))
+    result = processor.process(image, overlap=int(overlap))
     elapsed = time.time() - start
     
     result_pil = Image.fromarray(result)
@@ -310,19 +351,15 @@ with gr.Blocks(title="E3Diff: SAR-to-Optical Translation") as demo:
     
     **CVPR PBVS2025 Challenge Winner** | Upload any SAR image and get a photorealistic optical translation.
     
-    - Supports full resolution processing with seamless tiling
-    - Multiple quality levels (1-8 inference steps)
+    - Full resolution processing with seamless tiling
+    - One-step diffusion (optimized for speed & quality)
     - TIFF output for commercial use
     """)
     
     with gr.Row():
         with gr.Column():
             input_file = gr.File(label="SAR Input (TIFF, PNG, JPG)", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
-            
-            with gr.Row():
-                num_steps = gr.Slider(1, 8, value=1, step=1, label="Quality Steps (1=fast, 8=best)")
-                overlap = gr.Slider(16, 128, value=64, step=16, label="Tile Overlap")
-            
+            overlap = gr.Slider(16, 128, value=64, step=16, label="Tile Overlap (higher=smoother)")
             enhance = gr.Checkbox(value=True, label="Apply enhancement")
             submit_btn = gr.Button("🚀 Translate to Optical", variant="primary")
         
@@ -333,13 +370,13 @@ with gr.Blocks(title="E3Diff: SAR-to-Optical Translation") as demo:
     
     submit_btn.click(
         fn=translate_sar,
-        inputs=[input_file, num_steps, overlap, enhance],
+        inputs=[input_file, overlap, enhance],
         outputs=[output_image, output_file, info_text]
     )
     
     gr.Markdown("""
     ---
-    **Tips:** Use steps=1 for speed, steps=4-8 for quality. Works best with Sentinel-1 style SAR.
+    **Note:** E3Diff is a one-step diffusion model. Multiple steps degrade quality.
     """)