First Commit

app.py

@@ -0,0 +1,268 @@
+import gradio as gr
+import torch
+import time
+from PIL import Image, UnidentifiedImageError
+import os
+import requests
+import io
+import uuid
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+
+# Assuming these are correctly imported from your processing script
+from run_on_patches_online import (
+    load_model,
+    process_image_from_data,
+    Generator, # Make sure Generator is imported if needed by load_model
+    DEVICE,
+    #CHECKPOINT_GEN,
+    PATCH_KERNEL_SIZE,
+    PATCH_STRIDE
+)
+
+# --- Global Variables & Model Loading ---
+HF_REPO_ID = "b-aryan/WM-rem-epoch-42"
+HF_FILENAME = "gen_epoch_42.pth.tar"
+CHECKPOINT_GEN = HF_FILENAME
+MODEL = None
+MODEL_LOAD_ERROR = None
+DOWNLOADED_CHECKPOINT_PATH = None # Store the path after download
+
+print(f"Attempting to download/load model '{HF_FILENAME}' from repo '{HF_REPO_ID}' onto device '{DEVICE}'...")
+
+try:
+    # 1. Download the model checkpoint from Hugging Face Hub
+    print(f"Downloading checkpoint '{HF_FILENAME}' from '{HF_REPO_ID}'...")
+    DOWNLOADED_CHECKPOINT_PATH = hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=HF_FILENAME
+        # cache_dir can be specified if needed, otherwise uses default HF cache
+    )
+    print(f"Checkpoint downloaded successfully to: {DOWNLOADED_CHECKPOINT_PATH}")
+
+    # 2. Load the model using the downloaded path
+    if not os.path.exists(DOWNLOADED_CHECKPOINT_PATH):
+         # This should ideally not happen if hf_hub_download succeeded
+         raise FileNotFoundError(f"Downloaded checkpoint file not found at: {DOWNLOADED_CHECKPOINT_PATH}")
+
+    MODEL = load_model(DOWNLOADED_CHECKPOINT_PATH, DEVICE)
+    print("Model loaded successfully for Gradio app.")
+
+except Exception as e:
+    MODEL_LOAD_ERROR = f"Failed to download or load model '{HF_FILENAME}' from '{HF_REPO_ID}'. Error: {e}"
+    print(f"Error: {MODEL_LOAD_ERROR}")
+    import traceback
+    traceback.print_exc()
+
+TEMP_DIR = Path("temp")
+TEMP_DIR.mkdir(exist_ok=True)
+
+# --- Helper Function: Download Image (Simplified from run_on_patches_online) ---
+def download_image_for_gradio(url: str, timeout: int = 20) -> Image.Image | None:
+    """Downloads an image from a URL for Gradio, returns PIL Image or raises gr.Error."""
+    print(f"Attempting to download image from: {url}")
+    if not url or not url.startswith(('http://', 'https://')):
+        raise gr.Error("Please enter a valid HTTP or HTTPS image URL.")
+
+    try:
+        headers = {'User-Agent': 'Gradio-Image-Processor/1.1'}
+        response = requests.get(url, stream=True, timeout=timeout, headers=headers)
+        response.raise_for_status()
+
+        content_type = response.headers.get('Content-Type', '').lower()
+        if not content_type.startswith('image/'):
+            raise gr.Error(f"URL content type ({content_type}) is not recognized as an image.")
+
+        # Limit download size (e.g., 20 MB) to prevent abuse
+        content_length = response.headers.get('Content-Length')
+        if content_length and int(content_length) > 20 * 1024 * 1024:
+             raise gr.Error(f"Image file size ({int(content_length)/1024/1024:.1f} MB) exceeds the 20 MB limit.")
+
+        image_bytes = response.content
+        pil_image = Image.open(io.BytesIO(image_bytes))
+        pil_image = pil_image.convert('RGB')
+        print(f"Image downloaded successfully ({pil_image.width}x{pil_image.height}).")
+
+        # Optional: Add image dimension limits if needed
+        # max_dim = 2048
+        # if pil_image.width > max_dim or pil_image.height > max_dim:
+        #     raise gr.Error(f"Image dimensions ({pil_image.width}x{pil_image.height}) exceed the maximum allowed ({max_dim}x{max_dim}).")
+
+        return pil_image
+
+    except requests.exceptions.Timeout:
+        raise gr.Error(f"Request timed out after {timeout} seconds trying to download the image.")
+    except requests.exceptions.RequestException as e:
+        raise gr.Error(f"Error downloading image: {e}")
+    except UnidentifiedImageError:
+        raise gr.Error("Could not identify image file. The URL might not point to a valid image.")
+    except Exception as e:
+        print(f"An unexpected error occurred during download: {e}") # Log for server admin
+        raise gr.Error(f"An unexpected error occurred during image download.")
+
+
+# --- Processing Function (Handles the ML part with progress) ---
+def run_processing(input_pil_image: Image.Image, progress=gr.Progress(track_tqdm=True)):
+    """Processes the already downloaded PIL image and returns the result."""
+    if MODEL is None:
+        # Include the more specific error message if loading failed
+        error_msg = f"Model is not loaded. Cannot process image. Load error: {MODEL_LOAD_ERROR}" if MODEL_LOAD_ERROR else "Model is not loaded. Cannot process image."
+        raise gr.Error(error_msg)
+    if input_pil_image is None:
+        return None, None
+
+    start_time = time.time()
+    print("Starting image processing...")
+    progress(0, desc="Preparing for processing...")
+
+    try:
+        output_pil_image = process_image_from_data(
+            input_pil_image=input_pil_image,
+            model=MODEL,
+            device=DEVICE,
+            kernel_size=PATCH_KERNEL_SIZE,
+            stride=PATCH_STRIDE,
+            use_tqdm=True
+        )
+        if output_pil_image is None:
+            raise gr.Error("Image processing failed internally. Check server logs.")
+
+        # Save the processed image with a unique filename
+        filename = f"processed_{uuid.uuid4().hex}.jpg"
+        save_path = TEMP_DIR / filename
+        output_pil_image.save(save_path)
+        print(f"Saved processed image to: {save_path}")
+
+    except Exception as e:
+        print(f"Exception during processing: {e}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"An error occurred during model processing: {e}")
+
+    end_time = time.time()
+    processing_time = end_time - start_time
+    print(f"Processing time: {processing_time:.2f} seconds")
+
+    return output_pil_image, str(save_path)
+
+# --- Wrapper Function for Button Click ---
+def download_and_clear(url):
+    """Calls download, returns (image, None) on success, raises gr.Error on failure."""
+    print("Download and clear triggered.") # Debug print
+    try:
+        img = download_image_for_gradio(url)
+        # Return the downloaded image for the original slot, and None for the processed slot
+        return img, None
+    except gr.Error as e:
+        # If download fails, raise the Gradio error again
+        # Gradio will handle clearing outputs and displaying the error message
+        print(f"Download failed: {e}") # Debug print
+        raise e
+    except Exception as e:
+        # Catch other unexpected errors during the wrapper logic
+        print(f"Unexpected error in download_and_clear: {e}") # Debug print
+        import traceback
+        traceback.print_exc()
+        # Raise a Gradio error so the user sees something meaningful
+        raise gr.Error(f"An unexpected error occurred preparing the image: {e}")
+
+
+# --- Gradio Blocks Interface Definition ---
+print("Setting up Gradio Blocks interface...")
+
+with gr.Blocks() as demo:
+    saved_image_path = gr.State()
+    gr.Markdown(f"""
+        # Watermark Remover
+
+        **Disclaimer: This project was made to showcase my Deep Learning skills with no intention to cause harm to any business or infringe on any IP and it will be decommissioned as soon as I get a decent job offer. 
+        If you still have any issue, please reach out to me at [[email protected]](mailto:[email protected]).**
+    """)
+
+    gr.Markdown(f"""
+        This is a demo of a DL model which takes in URL of an image with watermark and gives out the image with its watermark removed.
+        The image is broken into overlapping patches of 256x256 pixels with a stride of 64 before feeding them to the model, 
+        the model infers on each patch separately and then they are stitched together to form the whole image again to give the final output.
+    """)
+
+    if MODEL_LOAD_ERROR:
+        gr.Markdown(f"<font color='red'>**Model Loading Error:** {MODEL_LOAD_ERROR}</font>")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_url = gr.Textbox(
+                label="Image URL",
+                placeholder="Paste image URL (https://www.shutterstock.com/shutterstock/photos/...)",
+                info="Enter the direct link to a publicly accessible image file (jpg, png, etc.)."
+            )
+            submit_button = gr.Button("Process Image", variant="primary")
+            gr.Markdown("---") # Separator
+            gr.Examples(
+                 examples=[
+                    ["https://www.shutterstock.com/shutterstock/photos/2429728793/display_1500/stock-photo-monkey-funny-open-mouth-and-showing-teeth-crazy-nature-angry-short-hair-brown-grand-bassin-2429728793.jpg"],
+                    ["https://www.shutterstock.com/shutterstock/photos/1905929728/display_1500/stock-photo-skeptic-surprised-cat-thinking-dont-know-what-to-do-big-eyes-closeup-tabby-cat-look-side-dont-1905929728.jpg"],
+                    ["https://www.shutterstock.com/shutterstock/photos/2501926843/display_1500/stock-photo-brunette-woman-laying-on-couch-cuddling-light-brown-dog-and-brown-tabby-cat-happy-2501926843.jpg"]
+                 ],
+                 inputs=image_url # Link examples to the URL textbox
+            )
+
+        with gr.Column(scale=2):
+            with gr.Row():
+                original_image = gr.Image(label="Original Input Image", type="pil", interactive=False)
+                processed_image = gr.Image(label="Processed Output Image", type="pil", interactive=False)
+            download_button = gr.DownloadButton("Download Processed Image", visible=True, variant="secondary")
+            gr.Markdown("---") # Separator
+            gr.Markdown(f"""
+                **A bit about the model:** In this project, I have trained a GAN network,
+                with the Generator being inspired from Pix2Pix and Pix2PixHD architectures and the Discriminator is very similar to PatchGAN in Pix2Pix and CycleGAN. 
+                For the loss, I have also added Perceptual Loss using VGG like in Pix2PixHD and SRGAN papers apart from the L1 and BCE loss.
+            """)
+
+            gr.Markdown(f"""If you liked this project, you can find my CV [here]() or reach me out at [[email protected]](mailto:[email protected]).""")
+
+
+    # --- Event Handling Logic ---
+
+    # 1. When Button is clicked:
+    #    - Input: URL from textbox
+    #    - Action: Call 'download_and_clear' wrapper function
+    #    - Output: Update 'original_image' with downloaded image, clear 'processed_image' by returning None for it.
+    submit_button.click(
+        fn=download_and_clear,  # <--- CORRECTED: Use the wrapper function
+        inputs=image_url,
+        outputs=[original_image, processed_image] # Target both outputs
+    )
+
+    # 2. When 'original_image' component *changes* (i.e., after successful download):
+    #    - Input: The PIL image data from 'original_image'
+    #    - Action: Call 'run_processing' (this function has the progress bar)
+    #    - Output: Update the 'processed_image' component.
+    original_image.change(
+        fn=run_processing,
+        inputs=original_image,
+        outputs=[processed_image, saved_image_path]
+        # concurrency_limit=1 # Optional: Prevent multiple simultaneous processing runs if needed
+    )
+
+    processed_image.change(
+        fn=lambda img: gr.update(variant="primary" if img is not None else "secondary"),
+        inputs=processed_image,
+        outputs=download_button
+    )
+
+    download_button.click(
+        fn=lambda path: path if path else None,
+        inputs=saved_image_path,
+        outputs=download_button
+    )
+
+
+
+# --- Launch the Application ---
+if __name__ == "__main__":
+    print("Launching Gradio Blocks interface...")
+    # Set queue=True for better handling under load, especially with long-running processing
+    # demo.queue() # Consider adding this
+    demo.launch(share=False, server_name="0.0.0.0", show_api=False)
+

generator_model.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, use_act=True, **kwargs):
+        super().__init__()
+        self.cnn = nn.Conv2d(in_channels, out_channels, **kwargs, bias=False, padding_mode="reflect")
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = nn.ReLU(inplace=True) if use_act else nn.Identity()
+
+    def forward(self, x):
+        return self.act(self.bn(self.cnn(x)))
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.survival_prob = 0.8
+        self.block1 = ConvBlock(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        self.block2 = ConvBlock(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            use_act=True,
+        )
+
+    def stochastic_depth(self, x):
+        if not self.training:
+            return x
+        binary_tensor = torch.rand(x.shape[0], 1, 1, 1, device=x.device) < self.survival_prob
+        return torch.div(x, self.survival_prob) * binary_tensor
+
+    def forward(self, x):
+        out = self.block1(x)
+        out = self.block2(out)
+        return self.stochastic_depth(out) + x
+
+class Block(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=2, act="relu"):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False, padding_mode="reflect"),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True) if act == "relu" else nn.LeakyReLU(0.2, inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Generator(nn.Module):
+    def __init__(self, in_channels=3, features=64, num_residuals=9):
+        super().__init__()
+        self.initial_down = nn.Sequential(
+            nn.Conv2d(in_channels, features, 7, 1, 3, bias=True, padding_mode="reflect"),
+            nn.ReLU(inplace=True),
+        )
+        self.down1 = Block(features, features*2, act="relu")
+        self.down2 = Block(features*2, features*4, act="relu")
+        self.down3 = Block(features*4, features*8, act="relu")
+        self.down4 = Block(features*8, features*16, act="relu")
+        self.residuals = nn.Sequential(*[ResidualBlock(features*16) for _ in range(num_residuals)])
+        self.up1 = Block(features*16, features*8, stride=1, act="relu")
+        self.up2 = Block(features*8*2, features*4, stride=1, act="relu" )
+        self.up3 = Block(features*4*2, features*2, stride=1, act="relu")
+        self.up4 = Block(features*2*2, features, stride=1, act="relu")
+        self.final_conv = nn.Sequential(
+            Block(features*2, features, stride=1, act="relu"),
+            Block(features, features, stride=1, act="relu"),
+            nn.Conv2d(features, in_channels, 7,1,3, padding_mode="reflect"),
+            nn.Tanh(),
+        )
+
+    def forward(self, x):
+        d1 = self.initial_down(x)
+        d2 = self.down1(d1)
+        d3 = self.down2(d2)
+        d4 = self.down3(d3)
+        d5 = self.down4(d4)
+        residuals = self.residuals(d5) + d5
+        u1 = self.up1(F.interpolate(residuals, scale_factor=2, mode="nearest"))
+        u2 = self.up2(F.interpolate(torch.cat([u1, d4], dim=1), scale_factor=2, mode="nearest"))
+        u3 = self.up3(F.interpolate(torch.cat([u2, d3], dim=1), scale_factor=2, mode="nearest"))
+        u4 = self.up4(F.interpolate(torch.cat([u3, d2], dim=1), scale_factor=2, mode="nearest"))
+        return self.final_conv(torch.cat([u4, d1], dim=1))
+
+def test():
+    x = torch.randn((1, 3, 256, 256))
+    model = Generator(in_channels=3, features=64)
+    preds = model(x)
+    print(preds.shape)
+
+#test()

requirements.txt

@@ -0,0 +1,8 @@
+gradio~=5.25.2
+requests~=2.32.3
+pillow~=11.1.0
+numpy~=1.26.4
+albumentations~=2.0.5
+tqdm~=4.67.1
+torch~=2.6.0
+huggingface_hub~=0.30.2

run_on_patches_online.py

@@ -0,0 +1,258 @@
+import os
+from PIL import Image, UnidentifiedImageError
+import numpy as np
+import torch
+import torch.nn.functional as F
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from collections import OrderedDict
+from tqdm import tqdm
+import requests
+import io
+from generator_model import Generator
+
+# --- Constants ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+CHECKPOINT_GEN = "gen_epoch_42.pth.tar" # Keep your checkpoint name
+PATCH_KERNEL_SIZE = 256
+PATCH_STRIDE = 64
+# DEFAULT_INPUT_DIR = "test/inputs"   # No longer needed for Gradio URL input
+# DEFAULT_OUTPUT_DIR = "test/outputs" # Output handled by Gradio
+SUPPORTED_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff') # Still useful for local testing if needed
+
+test_transform = A.Compose([
+    A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0),
+    ToTensorV2()
+])
+
+def load_model(checkpoint_path: str, device: str) -> Generator:
+    print(f"Loading model from: {checkpoint_path} onto device: {device}")
+    model = Generator(in_channels=3, features=64).to(device)
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    new_state_dict = OrderedDict()
+    has_module_prefix = any(k.startswith("module.") for k in checkpoint["state_dict"])
+    for k, v in checkpoint["state_dict"].items():
+        name = k.replace("module.", "") if has_module_prefix else k
+        new_state_dict[name] = v
+    model.load_state_dict(new_state_dict)
+    model.eval()  # Set model to evaluation mode
+    print("Model loaded successfully.")
+    return model
+
+def calculate_padding(img_h: int, img_w: int, kernel_size: int, stride: int) -> tuple[int, int]:
+    pad_h = kernel_size - img_h if img_h < kernel_size else (stride - (img_h - kernel_size) % stride) % stride
+    pad_w = kernel_size - img_w if img_w < kernel_size else (stride - (img_w - kernel_size) % stride) % stride
+    return pad_h, pad_w
+
+def download_image(url: str, timeout: int = 15) -> Image.Image | None:
+    """Downloads an image from a URL and returns it as a PIL Image object."""
+    print(f"Attempting to download image from: {url}")
+    try:
+        headers = {'User-Agent': 'Gradio-Image-Processor/1.0'} # Be a good net citizen
+        response = requests.get(url, stream=True, timeout=timeout, headers=headers)
+        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+
+        content_type = response.headers.get('Content-Type', '').lower()
+        if not content_type.startswith('image/'):
+            print(f"Error: URL content type ({content_type}) is not an image.")
+            return None
+
+        image_bytes = response.content
+        pil_image = Image.open(io.BytesIO(image_bytes))
+        pil_image = pil_image.convert('RGB') # Ensure image is in RGB format
+        print(f"Image downloaded successfully ({pil_image.width}x{pil_image.height}).")
+        return pil_image
+
+    except requests.exceptions.Timeout:
+        print(f"Error: Request timed out after {timeout} seconds.")
+        return None
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading image: {e}")
+        return None
+    except UnidentifiedImageError:
+        print("Error: Could not identify image file. The URL might not point to a valid image.")
+        return None
+    except Exception as e:
+        print(f"An unexpected error occurred during download: {e}")
+        return None
+
+def process_image_from_data(
+    input_pil_image: Image.Image,
+    model: Generator,
+    device: str,
+    kernel_size: int,
+    stride: int,
+    use_tqdm: bool = True # Optional: Control progress bar visibility
+    ) -> Image.Image | None:
+    """
+    Processes an input PIL image using the patch-based method and returns the output PIL image.
+    Returns None if an error occurs during processing.
+    """
+    print(f"\nProcessing image data...")
+    try:
+        image_np = np.array(input_pil_image) # Convert PIL Image to NumPy array
+        H, W, _ = image_np.shape
+        print(f"  Input dimensions: {W}x{H}")
+
+        # Apply transformations
+        transformed = test_transform(image=image_np)
+        input_tensor = transformed['image'].to(device)  # Shape: (C, H, W)
+        C = input_tensor.shape[0]
+
+        # Calculate and apply padding
+        pad_h, pad_w = calculate_padding(H, W, kernel_size, stride)
+        print(f"  Calculated padding (H, W): ({pad_h}, {pad_w})")
+        padded_tensor = F.pad(input_tensor.unsqueeze(0), (0, pad_w, 0, pad_h), mode='reflect').squeeze(0)
+        _, H_pad, W_pad = padded_tensor.shape
+        print(f"  Padded dimensions: {W_pad}x{H_pad}")
+
+        # Extract patches
+        patches = padded_tensor.unfold(1, kernel_size, stride).unfold(2, kernel_size, stride)
+        num_patches_h = patches.shape[1]
+        num_patches_w = patches.shape[2]
+        num_patches_total = num_patches_h * num_patches_w
+        print(f"  Extracted {num_patches_total} patches ({num_patches_h} H x {num_patches_w} W)")
+
+        patches = patches.contiguous().view(C, -1, kernel_size, kernel_size)
+        # Permute to (num_patches_total, C, kernel_size, kernel_size)
+        patches = patches.permute(1, 0, 2, 3).contiguous()
+
+        output_patches = []
+        # Set up tqdm iterator if enabled
+        patch_iterator = tqdm(patches, total=num_patches_total, desc="  Inferring patches", unit="patch", leave=False, disable=not use_tqdm)
+
+        # --- Inference Loop ---
+        with torch.no_grad():
+            for patch in patch_iterator:
+                # Add batch dimension, run model, remove batch dimension
+                output_patch = model(patch.unsqueeze(0)).squeeze(0)
+                # Move to CPU immediately to save GPU memory during inference loop
+                output_patches.append(output_patch.cpu())
+
+        # Stack output patches back together
+        # If GPU memory allows, move back for reconstruction, otherwise keep on CPU
+        # Let's try moving back to device for faster reconstruction if possible
+        try:
+            output_patches = torch.stack(output_patches).to(device)
+            print(f"  Output patches moved to {device} for reconstruction.")
+        except Exception as e: # Catch potential OOM on device
+            print(f"  Warning: Could not move all output patches to {device} ({e}). Reconstruction might be slower on CPU.")
+            output_patches = torch.stack(output_patches) # Keep on CPU
+
+
+        # --- Reconstruction ---
+        # Generate 2D Hann window for blending
+        window_1d = torch.hann_window(kernel_size, periodic=False, device=device) # periodic=False is common
+        window_2d = torch.outer(window_1d, window_1d)
+        window_2d = window_2d.unsqueeze(0).to(device) # Add channel dim and ensure on device
+
+        # Initialize output tensor and weight tensor (for weighted averaging)
+        output_tensor = torch.zeros((C, H_pad, W_pad), device=device, dtype=output_patches.dtype)
+        weight_tensor = torch.zeros((C, H_pad, W_pad), device=device, dtype=window_2d.dtype)
+
+        patch_idx = 0
+        reconstruct_iterator = tqdm(total=num_patches_total, desc="  Reconstructing", unit="patch", leave=False, disable=not use_tqdm)
+
+        for i in range(num_patches_h):
+            for j in range(num_patches_w):
+                h_start = i * stride
+                w_start = j * stride
+                h_end = h_start + kernel_size
+                w_end = w_start + kernel_size
+
+                # Get current patch (ensure it's on the correct device)
+                current_patch = output_patches[patch_idx].to(device)
+                weighted_patch = current_patch * window_2d # Apply window
+
+                # Add weighted patch to output tensor
+                output_tensor[:, h_start:h_end, w_start:w_end] += weighted_patch
+                # Accumulate weights
+                weight_tensor[:, h_start:h_end, w_start:w_end] += window_2d
+
+                patch_idx += 1
+                reconstruct_iterator.update(1)
+
+        reconstruct_iterator.close() # Close the inner tqdm bar
+
+        # Perform weighted averaging - clamp weights to avoid division by zero
+        output_averaged = output_tensor / weight_tensor.clamp(min=1e-6)
+
+        # Crop to original dimensions
+        output_cropped = output_averaged[:, :H, :W]
+        print(f"  Final output dimensions: {output_cropped.shape[2]}x{output_cropped.shape[1]}")
+
+        # --- Convert to Output Format ---
+        # Permute C, H, W -> H, W, C ; Move to CPU ; Convert to NumPy
+        output_numpy = output_cropped.permute(1, 2, 0).cpu().numpy()
+
+        # Denormalize: Assuming input was normalized to [-1, 1]
+        output_numpy = (output_numpy * 0.5 + 0.5) * 255.0
+
+        # Clip values to [0, 255] and convert to uint8
+        output_numpy = output_numpy.clip(0, 255).astype(np.uint8)
+
+        # Convert NumPy array back to PIL Image
+        output_image = Image.fromarray(output_numpy)
+
+        print("  Image processing complete.")
+        return output_image
+
+    except Exception as e:
+        print(f"Error during image processing: {e}")
+        import traceback
+        traceback.print_exc() # Print detailed traceback for debugging
+        return None
+
+if __name__ == "__main__":
+    print("--- Testing Phase 1 Refactoring ---")
+    print(f"Using device: {DEVICE}")
+    print(f"Using patch kernel size: {PATCH_KERNEL_SIZE}")
+    print(f"Using patch stride: {PATCH_STRIDE}")
+    print(f"Using model checkpoint: {CHECKPOINT_GEN}")
+
+    # 1. Load the model (as it would be done globally in Gradio app)
+    try:
+        model = load_model(CHECKPOINT_GEN, DEVICE)
+    except Exception as e:
+        print(f"Failed to load model. Exiting test. Error: {e}")
+        exit()
+
+
+    # 2. Test URL download
+    # Replace with a valid image URL for testing
+    # test_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"
+    test_url = "https://www.shutterstock.com/shutterstock/photos/2501926843/display_1500/stock-photo-brunette-woman-laying-on-couch-cuddling-light-brown-dog-and-brown-tabby-cat-happy-2501926843.jpg" # A smaller known image
+    input_pil = download_image(test_url)
+
+    if input_pil:
+        print(f"\nDownloaded image type: {type(input_pil)}, size: {input_pil.size}")
+
+        # 3. Test processing the downloaded image
+        output_pil = process_image_from_data(
+            input_pil_image=input_pil,
+            model=model,
+            device=DEVICE,
+            kernel_size=PATCH_KERNEL_SIZE,
+            stride=PATCH_STRIDE,
+            use_tqdm=True # Show progress bars during test
+        )
+
+        if output_pil:
+            print(f"\nProcessed image type: {type(output_pil)}, size: {output_pil.size}")
+            # Save the output locally for verification during testing
+            try:
+                os.makedirs("test_outputs", exist_ok=True)
+                output_filename = "test_output_" + os.path.basename(test_url).split('?')[0] # Basic filename extraction
+                if not output_filename.lower().endswith(SUPPORTED_EXTENSIONS):
+                     output_filename += ".png" # Ensure it has an extension
+                save_path = os.path.join("test_outputs", output_filename)
+                output_pil.save(save_path)
+                print(f"Saved test output to: {save_path}")
+            except Exception as e:
+                print(f"Error saving test output: {e}")
+        else:
+            print("\nImage processing failed.")
+    else:
+        print("\nImage download failed.")
+
+    print("\n--- Phase 1 Testing Complete ---")