Upload 2 files

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
 import torch
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
-from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 import gc
 from PIL import Image
 import numpy as np
@@ -22,8 +20,6 @@ class Config:
     MODEL_DIR = "weights"
     REALESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
     REALESRGAN_FILENAME = "RealESRGAN_x2plus.pth"
-    SWIN2SR_ID = "caidas/swin2SR-classical-sr-x2-64"
-    SD_ID = "stabilityai/stable-diffusion-x4-upscaler"
     
     # SOTA Models (2025)
     SPAN_URL = "https://huggingface.co/Phips/2xNomosUni_span_multijpg/resolve/main/2xNomosUni_span_multijpg.safetensors"
@@ -31,7 +27,6 @@ class Config:
     HATS_URL = "https://huggingface.co/Phips/4xNomos8kSCHAT-S/resolve/main/4xNomos8kSCHAT-S.safetensors"
     HATS_FILENAME = "4xNomos8kSCHAT-S.safetensors"
 
-    MAX_IMAGE_SIZE_SD = 512  # Max dimension for SD input to prevent OOM
     DEVICE = "cpu" # Force CPU for this demo, can be "cuda" if available
 
     @staticmethod
@@ -193,22 +188,22 @@ class RealESRGANStrategy(UpscalerStrategy):
                         # 'reduce-overhead' uses CUDA graphs, so only use it on CUDA
                         if Config.DEVICE == 'cuda':
                             self.model = torch.compile(self.model, mode='reduce-overhead')
-                            logger.info("✓ torch.compile enabled (reduce-overhead mode)")
+                            logger.info("[INFO] torch.compile enabled (reduce-overhead mode)")
                         elif os.name == 'nt' and Config.DEVICE == 'cpu':
                              # Windows requires MSVC for Inductor (default cpu backend)
                              # We skip it to avoid "Compiler: cl is not found" error unless user has it.
-                             logger.info("ℹ Skipping torch.compile on Windows CPU to avoid MSVC requirement.")
+                             logger.info("[INFO] Skipping torch.compile on Windows CPU to avoid MSVC requirement.")
                         elif (psutil.cpu_count(logical=False) or 0) < 4 and Config.DEVICE == 'cpu':
                              # Skip compilation on weak CPUs (e.g. HF Spaces Free Tier) to avoid long startup times
-                             logger.info("ℹ Skipping torch.compile on low-core CPU to prevent timeout.")
+                             logger.info("[INFO] Skipping torch.compile on low-core CPU to prevent timeout.")
                         else:
                             # On Linux/Mac CPU, use default mode or skip if problematic. Default is usually safe.
                             self.model = torch.compile(self.model)
-                            logger.info("✓ torch.compile enabled (default mode)")
+                            logger.info("[SUCCESS] torch.compile enabled (default mode)")
                             
                         self.compiled = True
                     except Exception as e:
-                        logger.warning(f"⚠ torch.compile not available or failed: {e}")
+                        logger.warning(f"[WARNING] torch.compile not available or failed: {e}")
                         self.compiled = True # Mark as tried
 
                 logger.info(f"{self.name} loaded successfully.")
@@ -270,102 +265,6 @@ class RealESRGANStrategy(UpscalerStrategy):
 
         return Image.fromarray(output_np)
 
-class Swin2SRStrategy(UpscalerStrategy):
-    def __init__(self):
-        super().__init__()
-        self.name = "Swin2SR x2"
-        self.processor = None
-
-    def load(self) -> None:
-        if self.model is None:
-            logger.info(f"Loading {self.name}...")
-            try:
-                self.processor = AutoImageProcessor.from_pretrained(Config.SWIN2SR_ID)
-                model = Swin2SRForImageSuperResolution.from_pretrained(Config.SWIN2SR_ID)
-                self.model = model.to(Config.DEVICE) # type: ignore
-                logger.info(f"{self.name} loaded successfully.")
-            except Exception as e:
-                logger.error(f"Failed to load Swin2SR: {e}")
-                # Swin2SR loading failure is often due to transformers version mismatch or device issues
-                # We re-raise to let the UI handle it, but log the specific error
-                raise
-
-    def upscale(self, image: Image.Image, **kwargs) -> Image.Image:
-        if self.model is None or self.processor is None:
-            self.load()
-            
-        logger.info(f"Starting inference with {self.name}...")
-        start_time = time.time()
-
-        if self.processor is None:
-             raise ValueError("Processor not loaded")
-
-        inputs = self.processor(images=image, return_tensors="pt").to(Config.DEVICE)
-        
-        # Swin2SR on CPU can be finicky with autocast/tracing. 
-        # Explicitly disable autocast for Swin2SR on CPU to avoid "PythonFallbackKernel" errors
-        context = torch.no_grad()
-        
-        with context:
-            outputs = self.model(**inputs)
-        
-        output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = np.moveaxis(output, source=0, destination=-1)
-        output = (output * 255.0).round().astype(np.uint8)
-        
-        logger.info(f"Inference finished in {time.time() - start_time:.2f}s")
-        return Image.fromarray(output)
-
-class StableDiffusionStrategy(UpscalerStrategy):
-    def __init__(self):
-        super().__init__()
-        self.name = "Stable Diffusion x4"
-
-    def load(self) -> None:
-        if self.model is None:
-            logger.info(f"Loading {self.name} (this may take time)...")
-            try:
-                self.model = StableDiffusionUpscalePipeline.from_pretrained(
-                    Config.SD_ID, 
-                    torch_dtype=torch.float32,
-                    low_cpu_mem_usage=True
-                )
-                # Optimizations for CPU
-                self.model.enable_attention_slicing("max")
-                self.model.enable_vae_tiling()
-                logger.info(f"{self.name} loaded successfully.")
-            except Exception as e:
-                logger.error(f"Failed to load Stable Diffusion: {e}")
-                raise
-
-    def upscale(self, image: Image.Image, **kwargs) -> Image.Image:
-        if self.model is None:
-            self.load()
-            
-        prompt = kwargs.get("prompt", "high quality, detailed")
-        
-        # Pre-check size
-        if max(image.size) > Config.MAX_IMAGE_SIZE_SD:
-            ratio = Config.MAX_IMAGE_SIZE_SD / max(image.size)
-            new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
-            image = image.resize(new_size, Image.Resampling.LANCZOS)
-            logger.warning(f"Resized input to {new_size} to prevent OOM on CPU.")
-
-        logger.info(f"Starting inference with {self.name}...")
-        start_time = time.time()
-        
-        generator = torch.manual_seed(42)
-        output = self.model(
-            prompt=prompt,
-            image=image,
-            num_inference_steps=20,
-            guidance_scale=7.0,
-            generator=generator
-        ).images[0] # type: ignore
-        
-        logger.info(f"Inference finished in {time.time() - start_time:.2f}s")
-        return output
-
 class SpanStrategy(UpscalerStrategy):
     def __init__(self):
         super().__init__()
@@ -401,14 +300,14 @@ class SpanStrategy(UpscalerStrategy):
                     try:
                         if Config.DEVICE == 'cuda':
                             self.model = torch.compile(self.model, mode='reduce-overhead')
-                            logger.info("✓ torch.compile enabled (reduce-overhead mode)")
+                            logger.info("[INFO] torch.compile enabled (reduce-overhead mode)")
                         elif os.name == 'nt' and Config.DEVICE == 'cpu':
-                             logger.info("ℹ Skipping torch.compile on Windows CPU.")
+                             logger.info("[INFO] Skipping torch.compile on Windows CPU.")
                         elif (psutil.cpu_count(logical=False) or 0) < 4 and Config.DEVICE == 'cpu':
-                             logger.info("ℹ Skipping torch.compile on low-core CPU.")
+                             logger.info("[INFO] Skipping torch.compile on low-core CPU.")
                         else:
                             # SPAN architecture uses .data.clone() in forward pass which breaks torch.compile/inductor
-                            logger.info("ℹ Skipping torch.compile for SPAN (incompatible architecture).")
+                            logger.info("[INFO] Skipping torch.compile for SPAN (incompatible architecture).")
                             # self.model = torch.compile(self.model) 
                         self.compiled = True
                     except Exception:
@@ -503,7 +402,7 @@ class HatsStrategy(UpscalerStrategy):
                              pass
                         else:
                             # HAT architecture also triggers "UntypedStorage" weakref errors with inductor on CPU
-                            logger.info("ℹ Skipping torch.compile for HAT-S (incompatible architecture).")
+                            logger.info("[INFO] Skipping torch.compile for HAT-S (incompatible architecture).")
                             # self.model = torch.compile(self.model)
                         self.compiled = True
                     except Exception:
@@ -560,9 +459,7 @@ class UpscalerManager:
         self.strategies: Dict[str, UpscalerStrategy] = {
             "SPAN (NomosUni) x2": SpanStrategy(),
             "RealESRGAN x2": RealESRGANStrategy(),
-            "HAT-S x4": HatsStrategy(),
-            "Swin2SR x2": Swin2SRStrategy(),
-            "Stable Diffusion x4": StableDiffusionStrategy()
+            "HAT-S x4": HatsStrategy()
         }
         self.current_model_name: Optional[str] = None
 
@@ -590,26 +487,31 @@ class UpscalerManager:
 manager = UpscalerManager()
 
 # --- Gradio Interface Logic ---
-def process_image(input_img: Image.Image, model_name: str, prompt: str) -> Tuple[Optional[Image.Image], str, str]:
+def process_image(input_img: Image.Image, model_name: str, output_format: str) -> Tuple[Optional[str], str, str]:
     if input_img is None:
         return None, get_logs(), get_system_usage()
     
     try:
         strategy = manager.get_strategy(model_name)
         
-        # Optional: Unload others if memory is tight (simple logic here)
-        # For now, we just rely on the user or OS, but in prod we might auto-unload.
+        output_img = strategy.upscale(input_img)
         
-        output = strategy.upscale(input_img, prompt=prompt)
+        # Save to temp file with correct extension
+        output_path = f"output.{output_format.lower()}"
+        
+        # Convert to RGB if saving as JPEG (doesn't support alpha)
+        if output_format.lower() in ['jpeg', 'jpg'] and output_img.mode == 'RGBA':
+            output_img = output_img.convert('RGB')
+            
+        output_img.save(output_path, format=output_format)
         
         # Explicit GC after heavy operations
         gc.collect()
         
-        return output, get_logs(), get_system_usage()
+        return output_path, get_logs(), get_system_usage()
     except Exception as e:
         error_msg = f"Critical Error: {str(e)}\n{traceback.format_exc()}"
         logger.error(error_msg)
-        # Return the error message in the logs output so the user sees it
         return None, get_logs() + "\n\n" + error_msg, get_system_usage()
 
 def unload_models():
@@ -618,49 +520,61 @@ def unload_models():
 
 # --- UI Construction ---
 desc = """
-### 🚀 Enterprise-Grade Universal Upscaler (SOTA 2025)
-Select a specialized model to upscale your image.
-* **SPAN (NomosUni) x2**: ⚡ **SOTA Speed**. Fastest CPU model. Best for general use.
-* **RealESRGAN x2**: 🛡️ **Robust**. Best for removing JPEG artifacts and noise.
-* **HAT-S x4**: 💎 **SOTA Quality**. Best texture detail (slower).
-* **Swin2SR x2**: 🎯 High fidelity, removes compression artifacts.
-* **Stable Diffusion x4**: 🎨 Generative upscaling. Adds missing details (slow, high RAM).
+# Universal Upscaler Pro (CPU Optimized)
+
+This application provides state-of-the-art (SOTA) image upscaling running entirely on CPU, optimized for free-tier cloud environments.
+
+### Available Models
+
+| Model | Scale | Best For | License |
+| :--- | :--- | :--- | :--- |
+| **SPAN (NomosUni)** | x2 | **Speed & General Use**. Extremely fast, parameter-free attention network. | Apache 2.0 |
+| **RealESRGAN** | x2 | **Robustness**. Excellent at removing JPEG artifacts and noise. | BSD 3-Clause |
+| **HAT-S** | x4 | **Texture Detail**. Hybrid Attention Transformer for high-fidelity restoration. | MIT |
+
+### Attributions & Credits
+
+*   **Real-ESRGAN**: [Wang et al., 2021](https://github.com/xinntao/Real-ESRGAN). *Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data*.
+*   **SPAN**: [Zhang et al., 2023](https://github.com/hongyuanyu/SPAN). *Swift Parameter-free Attention Network for Efficient Super-Resolution*.
+*   **HAT**: [Chen et al., 2023](https://github.com/XPixelGroup/HAT). *Activating Activation Functions for Image Restoration*.
+*   **NomosUni**: Custom SPAN training by [Phhofm](https://github.com/Phhofm).
 """
 
 with gr.Blocks(title="Universal Upscaler Pro") as iface:
     gr.Markdown(desc)
     
     with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="Input Image")
+        with gr.Column(scale=1, min_width=300):
+            input_image = gr.Image(type="pil", label="Input Image", height=400)
             
-            with gr.Group():
+            with gr.Row():
                 model_selector = gr.Dropdown(
                     choices=list(manager.strategies.keys()),
                     value="SPAN (NomosUni) x2",
-                    label="Select Model Architecture"
+                    label="Model Architecture",
+                    scale=2
                 )
-                prompt_input = gr.Textbox(
-                    label="Prompt (Stable Diffusion Only)", 
-                    value="highly detailed, 4k, sharp",
-                    placeholder="Describe the image content..."
+                output_format = gr.Dropdown(
+                    choices=["PNG", "JPEG", "WEBP"],
+                    value="PNG",
+                    label="Output Format",
+                    scale=1
                 )
             
+            submit_btn = gr.Button("Upscale Image", variant="primary", size="lg")
+            
             with gr.Accordion("Advanced Settings", open=False):
-                gr.Markdown("Memory Management")
                 unload_btn = gr.Button("Unload All Models (Free RAM)", variant="secondary")
-            
-            submit_btn = gr.Button("✨ Upscale Image", variant="primary", size="lg")
-            system_info = gr.Label(value=get_system_usage(), label="System Status")
+                system_info = gr.Label(value=get_system_usage(), label="System Status")
 
-        with gr.Column(scale=1):
-            output_image = gr.Image(type="pil", label="Upscaled Result")
-            logs_output = gr.TextArea(label="Execution Logs", interactive=False, lines=10)
+        with gr.Column(scale=1, min_width=300):
+            output_image = gr.Image(type="filepath", label="Upscaled Result", height=400)
+            logs_output = gr.TextArea(label="Execution Logs", interactive=False, lines=8)
 
     # Event Wiring
     submit_btn.click(
         fn=process_image,
-        inputs=[input_image, model_selector, prompt_input],
+        inputs=[input_image, model_selector, output_format],
         outputs=[output_image, logs_output, system_info]
     )
     
@@ -669,8 +583,5 @@ with gr.Blocks(title="Universal Upscaler Pro") as iface:
         inputs=[],
         outputs=[logs_output, system_info]
     )
-    
-    # Auto-refresh system info every 2 seconds (optional, can be heavy on UI)
-    # iface.load(get_system_usage, None, system_info, every=2)
 
-iface.launch()
+iface.launch()

requirements.txt

@@ -1,7 +1,4 @@
 torch
-diffusers
-transformers
-accelerate
 scipy
 pillow
 gradio
@@ -14,4 +11,4 @@ onnxruntime
 basicsr
 realesrgan
 openvino
-optimum
+optimum