Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,13 +7,14 @@ import io
 import numpy as np
 import gc
 
-# Page Configuration d
+# Page Configuration
 st.set_page_config(layout="wide", page_title="AI Image Lab")
 
 # --- 1. MODEL LOADING ---
 
 @st.cache_resource
 def load_rembg_model():
+    # RMBG-1.4 (Fast & High Quality)
     model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -22,9 +23,12 @@ def load_rembg_model():
 @st.cache_resource
 def load_upscaler(scale=2):
     if scale == 4:
-        model_id = "caidas/swin2SR-classical-sr-x4-63"
+        # FIXED: Use the 'RealWorld' model for 4x. It exists and handles artifacts better.
+        model_id = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr"
     else:
+        # 2x Classical Model
         model_id = "caidas/swin2SR-classical-sr-x2-64"
+        
     processor = AutoImageProcessor.from_pretrained(model_id)
     model = Swin2SRForImageSuperResolution.from_pretrained(model_id)
     return processor, model
@@ -32,6 +36,7 @@ def load_upscaler(scale=2):
 # --- 2. PROCESSING LOGIC ---
 
 def find_mask_tensor(output):
+    """Recursively finds the mask tensor."""
     if isinstance(output, torch.Tensor):
         if output.dim() == 4 and output.shape[1] == 1: return output
         elif output.dim() == 3 and output.shape[0] == 1: return output
@@ -64,7 +69,7 @@ def upscale_chunk_logic(image, processor, model):
 
 @st.cache_data(show_spinner=False)
 def process_background_removal(image_bytes):
-    """Cached background removal."""
+    """Cached background removal (RMBG-1.4)."""
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     model, device = load_rembg_model()
     
@@ -94,17 +99,16 @@ def process_background_removal(image_bytes):
 
 def process_tiled_upscale(image, scale_factor, grid_n, progress_bar):
     """
-    Tiled upscaling with OVERLAP to prevent edge artifacts.
+    Tiled upscaling with OVERLAP to prevent seams.
     """
     processor, model = load_upscaler(scale_factor)
     w, h = image.size
     rows = cols = grid_n
     
-    # Base tile size (without overlap)
     tile_w = w // cols
     tile_h = h // rows
     
-    # Overlap buffer (pixels) - lets the AI see context
+    # Overlap buffer (pixels)
     overlap = 32 
     
     full_image = Image.new(image.mode, (w * scale_factor, h * scale_factor))
@@ -113,52 +117,40 @@ def process_tiled_upscale(image, scale_factor, grid_n, progress_bar):
     
     for y in range(rows):
         for x in range(cols):
-            # 1. Define the "Target" area (where this tile goes in the original)
+            # Target Area
             target_left = x * tile_w
             target_upper = y * tile_h
-            # Handle edge pixels for the last column/row
             target_right = w if x == cols - 1 else (x + 1) * tile_w
             target_lower = h if y == rows - 1 else (y + 1) * tile_h
-            
             target_w = target_right - target_left
             target_h = target_lower - target_upper
 
-            # 2. Define the "Source" area (Target + Overlap)
-            # We expand the box outwards by 'overlap' px, but keep it within image bounds
+            # Source Area (with overlap)
             source_left = max(0, target_left - overlap)
             source_upper = max(0, target_upper - overlap)
             source_right = min(w, target_right + overlap)
             source_lower = min(h, target_lower + overlap)
             
-            # Crop the padded tile
             tile = image.crop((source_left, source_upper, source_right, source_lower))
-            
-            # 3. Upscale the Padded Tile
             upscaled_tile = upscale_chunk_logic(tile, processor, model)
             
-            # 4. Crop the "Valid" center from the upscaled tile
-            # Calculate how much extra we took on the Left and Top (in original scale)
+            # Calculate offsets for cropping the valid center
             extra_left = target_left - source_left
             extra_upper = target_upper - source_upper
             
-            # Convert these offsets to the new Upscaled Scale
             crop_x = extra_left * scale_factor
             crop_y = extra_upper * scale_factor
             crop_w = target_w * scale_factor
             crop_h = target_h * scale_factor
             
-            # Perform the final crop to remove the overlap borders
             clean_tile = upscaled_tile.crop((crop_x, crop_y, crop_x + crop_w, crop_y + crop_h))
             
-            # 5. Paste the clean tile
             paste_x = target_left * scale_factor
             paste_y = target_upper * scale_factor
             full_image.paste(clean_tile, (paste_x, paste_y))
             
-            # Cleanup
             del tile, upscaled_tile, clean_tile
             gc.collect()
-            
             count += 1
             progress_bar.progress(count / total_tiles, text=f"Upscaling Tile {count}/{total_tiles} (with overlap)...")
 
@@ -172,7 +164,7 @@ def convert_image_to_bytes(img):
 # --- 3. MAIN APP ---
 
 def main():
-    st.title("✨ AI Image Lab: Seamless Edition")
+    st.title("✨ AI Image Lab: Final Edition")
     st.markdown("Features: **RMBG-1.4** | **Swin2SR (Seamless Tiling)** | **Progress Bar**")
 
     # --- Sidebar ---
@@ -207,7 +199,7 @@ def main():
             scale = 4 if "4x" in upscale_mode else 2
             
             # Cache Key
-            cache_key = f"{uploaded_file.name}_{remove_bg}_{scale}_{grid_n}_overlap"
+            cache_key = f"{uploaded_file.name}_{remove_bg}_{scale}_{grid_n}_overlap_v4"
             
             if "upscale_cache" not in st.session_state:
                 st.session_state.upscale_cache = {}