Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,100 +3,166 @@ from PIL import Image, ImageEnhance
 import torch
 import torch.nn.functional as F
 from torchvision import transforms
-from transformers import AutoModelForImageSegmentation
+from transformers import AutoModelForImageSegmentation, AutoImageProcessor, Swin2SRForImageSuperResolution
 import io
 import numpy as np
 
 # Page Configuration
 st.set_page_config(layout="wide", page_title="AI Image Lab")
 
-# --- Caching AI Models ---
+# --- 1. MODEL LOADING (Cached) ---
 
 @st.cache_resource
-def load_birefnet_model():
-    """
-    Loads the RMBG-1.4 model for Background Removal (Pure PyTorch).
-    """
-    # We use 'briaai/RMBG-1.4' which is SOTA for background removal
+def load_rembg_model():
+    """Loads RMBG-1.4 for Background Removal."""
     model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
-    
-    # Move to GPU if available, otherwise CPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
     return model, device
 
-# ... (Previous Upscalers kept for reference, you can re-add them if you wish) ...
-
-def remove_background_torch(image, model, device):
+@st.cache_resource
+def load_upscaler(scale=2):
+    """Loads Swin2SR for Super-Resolution (2x or 4x)."""
+    if scale == 4:
+        model_id = "caidas/swin2SR-classical-sr-x4-63"
+    else:
+        model_id = "caidas/swin2SR-classical-sr-x2-64"
+        
+    processor = AutoImageProcessor.from_pretrained(model_id)
+    model = Swin2SRForImageSuperResolution.from_pretrained(model_id)
+    return processor, model
+
+# --- 2. PROCESSING FUNCTIONS ---
+
+def safe_rembg_inference(model, image, device):
     """
-    Runs background removal using RMBG-1.4 on PyTorch.
+    Robust inference for RMBG-1.4 that handles different output formats.
     """
-    # 1. Prepare input
     w, h = image.size
     
-    # The model expects specific normalization and size
+    # Preprocessing
     transform_image = transforms.Compose([
         transforms.Resize((1024, 1024)),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
-
     input_images = transform_image(image).unsqueeze(0).to(device)
 
-    # 2. Inference
+    # Inference
     with torch.no_grad():
-        preds = model(input_images)[-1].sigmoid().cpu()
+        outputs = model(input_images)
     
-    # 3. Post-process mask
-    pred = preds[0].squeeze()
+    # FIX: Handle List vs Tuple vs Tensor output
+    # BiRefNet usually returns a list/tuple of tensors. 
+    # The output we want is usually the LAST element or the FIRST depending on version.
+    # We check if 'outputs' is a sequence (list/tuple) and grab the tensor.
+    if isinstance(outputs, (list, tuple)):
+        # We assume the last element is the high-res prediction for RMBG-1.4
+        result_tensor = outputs[-1]
+        
+        # Double check: if the result is still a list (nested), grab the first item
+        if isinstance(result_tensor, (list, tuple)):
+            result_tensor = result_tensor[0]
+    else:
+        result_tensor = outputs
+
+    # Post-processing
+    pred = result_tensor.sigmoid().cpu()[0].squeeze()
     
-    # Convert mask to PIL and resize back to original dimensions
+    # Convert mask to PIL
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize((w, h))
     
-    # 4. Apply mask to original image
+    # Apply mask
     image.putalpha(mask)
     return image
 
+def ai_upscale(image, processor, model):
+    """
+    Upscales RGB image using Swin2SR. 
+    Note: Swin2SR only works on RGB. If image is RGBA, we must handle Alpha separately.
+    """
+    # 1. Handle Alpha Channel (if exists)
+    if image.mode == 'RGBA':
+        # Split RGB and Alpha
+        r, g, b, a = image.split()
+        rgb_image = Image.merge('RGB', (r, g, b))
+        
+        # Upscale RGB using AI
+        upscaled_rgb = run_swin_inference(rgb_image, processor, model)
+        
+        # Upscale Alpha using standard interpolation (AI models don't predict alpha)
+        # We resize alpha to match the new RGB size
+        upscaled_a = a.resize(upscaled_rgb.size, Image.Resampling.LANCZOS)
+        
+        # Recombine
+        return Image.merge('RGBA', (*upscaled_rgb.split(), upscaled_a))
+    else:
+        return run_swin_inference(image, processor, model)
+
+def run_swin_inference(image, processor, model):
+    """Helper to run the actual Swin2SR inference on an RGB image."""
+    inputs = processor(image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+    output = np.moveaxis(output, 0, -1)
+    output = (output * 255.0).round().astype(np.uint8)
+    return Image.fromarray(output)
+
 def convert_image_to_bytes(img):
     buf = io.BytesIO()
     img.save(buf, format="PNG")
     return buf.getvalue()
 
+# --- 3. MAIN APP ---
+
 def main():
-    st.title("✨ AI Image Lab: Pure PyTorch Edition")
-    st.markdown("Processing pipeline: **RMBG-1.4 (No ONNX)**")
-    st.write("XSRF:", st.get_option("server.enableXsrfProtection"))
+    st.title("✨ AI Image Lab: Robust Edition")
+    st.markdown("Features: **RMBG-1.4 (No ONNX)** | **Swin2SR (Upscaling)** | **Geometry**")
 
-    # --- Sidebar Controls ---
-    st.sidebar.header("Processing Pipeline")
+    # --- Sidebar ---
+    st.sidebar.header("1. Background")
+    remove_bg = st.sidebar.checkbox("Remove Background", value=False)
     
-    remove_bg = st.sidebar.checkbox("Remove Background (RMBG-1.4)", value=False)
+    st.sidebar.header("2. AI Upscaling")
+    upscale_mode = st.sidebar.radio("Magnification", ["None", "2x (Fast)", "4x (Slow)"])
     
-    st.sidebar.subheader("Final Adjustments")
+    st.sidebar.header("3. Geometry")
     rotate_angle = st.sidebar.slider("Rotate", -180, 180, 0, 1)
 
-    # --- Main Content ---
-    uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png", "webp"])
+    # --- Main ---
+    uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "webp"])
 
     if uploaded_file is not None:
-        # Important: RMBG model works best if we ensure RGB mode
         image = Image.open(uploaded_file).convert("RGB")
+        
+        # Create a working copy
         processed_image = image.copy()
-
-        # --- STEP 1: Background Removal ---
+        
+        # 1. Remove Background (Do this first so we have the mask)
         if remove_bg:
-            st.info("Loading RMBG-1.4 Model (First run will download ~170MB)...")
+            st.info("Loading RMBG Model...")
             try:
-                # Load Model
-                model, device = load_birefnet_model()
-                
-                with st.spinner("Removing background using PyTorch..."):
-                    processed_image = remove_background_torch(processed_image, model, device)
+                bg_model, device = load_rembg_model()
+                with st.spinner("Removing background..."):
+                    processed_image = safe_rembg_inference(bg_model, processed_image, device)
             except Exception as e:
-                st.error(f"Error during background removal: {e}")
+                st.error(f"Background Removal Failed: {e}")
 
-        # --- STEP 2: Geometry/Color ---
+        # 2. Upscaling
+        if upscale_mode != "None":
+            scale = 4 if "4x" in upscale_mode else 2
+            st.info(f"Loading Swin2SR x{scale} Model...")
+            try:
+                processor, upscaler = load_upscaler(scale)
+                with st.spinner(f"Upscaling x{scale}..."):
+                    processed_image = ai_upscale(processed_image, processor, upscaler)
+            except Exception as e:
+                st.error(f"Upscaling Failed: {e}")
+
+        # 3. Rotation
         if rotate_angle != 0:
             processed_image = processed_image.rotate(rotate_angle, expand=True)
 
@@ -105,18 +171,20 @@ def main():
         with col1:
             st.subheader("Original")
             st.image(image, use_container_width=True)
-
+            st.caption(f"Size: {image.size}")
+        
         with col2:
             st.subheader("Result")
             st.image(processed_image, use_container_width=True)
+            st.caption(f"Size: {processed_image.size}")
 
         # --- Download ---
         st.markdown("---")
-        btn = st.download_button(
-            label="💾 Download Result",
+        st.download_button(
+            label="💾 Download Result (PNG)",
             data=convert_image_to_bytes(processed_image),
-            file_name="nobg_image.png",
-            mime="image/png",
+            file_name="processed_image.png",
+            mime="image/png"
         )
 
 if __name__ == "__main__":