Upload app.py with huggingface_hub

app.py

@@ -10,17 +10,17 @@ import gradio as gr
 import tempfile
 import time
 
-print("[E3Diff] Starting app...")
+print("[Axion] Starting app...")
 
 # ZeroGPU support
 try:
     import spaces
     GPU_AVAILABLE = True
-    print("[E3Diff] ZeroGPU available")
+    print("[Axion] ZeroGPU available")
 except ImportError:
     GPU_AVAILABLE = False
     spaces = None
-    print("[E3Diff] Running without ZeroGPU")
+    print("[Axion] Running without ZeroGPU")
 
 
 # Lazy imports for heavy modules
@@ -30,20 +30,20 @@ _model_modules = None
 def get_torch():
     global _torch
     if _torch is None:
-        print("[E3Diff] Importing torch...")
+        print("[Axion] Importing torch...")
         import torch
         _torch = torch
-        print(f"[E3Diff] PyTorch {torch.__version__} loaded")
+        print(f"[Axion] PyTorch {torch.__version__} loaded")
     return _torch
 
 def get_model_modules():
     global _model_modules
     if _model_modules is None:
-        print("[E3Diff] Importing model modules...")
+        print("[Axion] Importing model modules...")
         from unet import UNet
         from diffusion import GaussianDiffusion
         _model_modules = (UNet, GaussianDiffusion)
-        print("[E3Diff] Model modules loaded")
+        print("[Axion] Model modules loaded")
     return _model_modules
 
 
@@ -87,7 +87,7 @@ def build_model(device):
     UNet, GaussianDiffusion = get_model_modules()
     from huggingface_hub import hf_hub_download
     
-    print("[E3Diff] Building model architecture...")
+    print("[Axion] Building model architecture...")
     
     image_size = 256
     num_inference_steps = 1
@@ -142,18 +142,18 @@ def build_model(device):
     model = model.to(device)
     
     # Load weights
-    print("[E3Diff] Downloading weights...")
+    print("[Axion] Downloading weights...")
     weights_path = hf_hub_download(
         repo_id="Dhenenjay/E3Diff-SAR2Optical",
         filename="I700000_E719_gen.pth"
     )
     
-    print(f"[E3Diff] Loading weights from: {weights_path}")
+    print(f"[Axion] Loading weights from: {weights_path}")
     state_dict = torch.load(weights_path, map_location=device, weights_only=False)
     model.load_state_dict(state_dict, strict=False)
     model.eval()
     
-    print("[E3Diff] Model ready!")
+    print("[Axion] Model ready!")
     return model
 
 
@@ -262,7 +262,7 @@ def process_image(image, model, device, overlap=64):
     x_positions = list(range(0, w_pad - tile_size + 1, step))
     total_tiles = len(y_positions) * len(x_positions)
     
-    print(f"[E3Diff] Processing {total_tiles} tiles ({len(x_positions)}x{len(y_positions)}) at {w}x{h}...")
+    print(f"[Axion] Processing {total_tiles} tiles ({len(x_positions)}x{len(y_positions)}) at {w}x{h}...")
     
     tile_idx = 0
     for y in y_positions:
@@ -281,7 +281,7 @@ def process_image(image, model, device, overlap=64):
             
             tile_idx += 1
             if tile_idx % 10 == 0 or tile_idx == total_tiles:
-                print(f"[E3Diff] Tile {tile_idx}/{total_tiles}")
+                print(f"[Axion] Tile {tile_idx}/{total_tiles}")
     
     # Normalize
     output = output / (weights + 1e-8)
@@ -303,7 +303,7 @@ def _translate_impl(file, overlap, enhance_output):
     
     torch = get_torch()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"[E3Diff] Using device: {device}")
+    print(f"[Axion] Using device: {device}")
     
     # Load model (cached)
     if _cached_model is None:
@@ -313,11 +313,11 @@ def _translate_impl(file, overlap, enhance_output):
     
     # Load image
     filepath = file.name if hasattr(file, 'name') else file
-    print(f"[E3Diff] Loading: {filepath}")
+    print(f"[Axion] Loading: {filepath}")
     image = load_sar_image(filepath)
     
     w, h = image.size
-    print(f"[E3Diff] Input size: {w}x{h}")
+    print(f"[Axion] Input size: {w}x{h}")
     
     start = time.time()
     result = process_image(image, model, device, overlap=int(overlap))
@@ -331,7 +331,7 @@ def _translate_impl(file, overlap, enhance_output):
     tiff_path = tempfile.mktemp(suffix='.tiff')
     result_pil.save(tiff_path, format='TIFF', compression='lzw')
     
-    print(f"[E3Diff] Complete in {elapsed:.1f}s!")
+    print(f"[Axion] Complete in {elapsed:.1f}s!")
     
     info = f"Processed in {elapsed:.1f}s | Output: {result_pil.size[0]}x{result_pil.size[1]}"
     
@@ -347,31 +347,58 @@ else:
     translate_sar = _translate_impl
 
 
-print("[E3Diff] Building Gradio interface...")
+print("[Axion] Building Gradio interface...")
+
+# Custom CSS for dark minimal theme
+custom_css = """
+.gradio-container {
+    background: linear-gradient(180deg, #0a0a0a 0%, #1a1a1a 100%) !important;
+}
+.main-title {
+    font-family: 'Helvetica Neue', Arial, sans-serif !important;
+    font-size: 3.5rem !important;
+    font-weight: 200 !important;
+    color: #ffffff !important;
+    text-align: center !important;
+    margin-bottom: 0.5rem !important;
+    letter-spacing: -0.02em !important;
+}
+.subtitle {
+    font-family: 'Helvetica Neue', Arial, sans-serif !important;
+    font-size: 1.1rem !important;
+    font-weight: 300 !important;
+    color: #888888 !important;
+    text-align: center !important;
+    margin-bottom: 2rem !important;
+}
+.dark-panel {
+    background: rgba(30, 30, 30, 0.6) !important;
+    border: 1px solid #333 !important;
+    border-radius: 12px !important;
+}
+"""
 
 # Create Gradio interface
-with gr.Blocks(title="E3Diff: SAR-to-Optical Translation") as demo:
-    gr.Markdown("""
-    # 🛰️ E3Diff: High-Resolution SAR-to-Optical Translation
-    
-    **CVPR PBVS2025 Challenge Winner** | Upload any SAR image and get a photorealistic optical translation.
-    
-    - Full resolution processing with seamless tiling
-    - One-step diffusion (optimized for speed & quality)
-    - TIFF output for commercial use
+with gr.Blocks(title="Axion - SAR to Optical", css=custom_css) as demo:
+    gr.HTML("""
+    <div style="text-align: center; padding: 40px 20px 20px 20px; background: linear-gradient(180deg, #0a0a0a 0%, #1a1a1a 100%);">
+        <h1 style="font-family: 'Helvetica Neue', Arial, sans-serif; font-size: 3.2rem; font-weight: 200; color: #ffffff; margin-bottom: 0.5rem; letter-spacing: -0.02em;">SAR to Optical Image Translation</h1>
+        <p style="font-family: 'Helvetica Neue', Arial, sans-serif; font-size: 1.1rem; font-weight: 300; color: #888888;">Transform radar imagery into crystal-clear optical views using our foundation model</p>
+    </div>
     """)
     
     with gr.Row():
         with gr.Column():
-            input_file = gr.File(label="SAR Input (TIFF, PNG, JPG)", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
-            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")
+            input_file = gr.File(label="Upload SAR Image", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
+            with gr.Row():
+                overlap = gr.Slider(16, 128, value=64, step=16, label="Tile Overlap")
+                enhance = gr.Checkbox(value=True, label="Enhance Output")
+            submit_btn = gr.Button("Translate", variant="primary")
         
         with gr.Column():
             output_image = gr.Image(label="Optical Output")
-            output_file = gr.File(label="Download TIFF")
-            info_text = gr.Textbox(label="Processing Info")
+            output_file = gr.File(label="Download")
+            info_text = gr.Textbox(label="Info", show_label=False)
     
     submit_btn.click(
         fn=translate_sar,
@@ -379,12 +406,13 @@ with gr.Blocks(title="E3Diff: SAR-to-Optical Translation") as demo:
         outputs=[output_image, output_file, info_text]
     )
     
-    gr.Markdown("""
-    ---
-    **Note:** E3Diff is a one-step diffusion model. Multiple steps degrade quality.
+    gr.HTML("""
+    <div style="text-align: center; padding: 20px; color: #555; font-size: 0.85rem;">
+        Powered by <strong style="color: #888;">Axion</strong>
+    </div>
     """)
 
-print("[E3Diff] Launching app...")
+print("[Axion] Launching app...")
 
 if __name__ == "__main__":
     demo.queue().launch(ssr_mode=False)