Upload folder using huggingface_hub

app.py

@@ -41,27 +41,56 @@ pipe = pipe.to(device)
 print("Model loaded successfully!")
 
 
-def preprocess_image(image: Image.Image) -> tuple[Image.Image, tuple[int, int]]:
-    """Pad image to multiple of 64 for SD3 compatibility."""
+def preprocess_image(image: Image.Image, max_size: int = 1536) -> tuple[Image.Image, tuple[int, int], float]:
+    """
+    Resize large images and pad to multiple of 64 for SD3 compatibility.
+    
+    Returns:
+        Tuple of (processed_image, original_size, scale_factor)
+    """
     original_size = image.size
     w, h = original_size
     
-    # SD3 works best with dimensions that are multiples of 64
-    new_w = (w + 63) // 64 * 64
-    new_h = (h + 63) // 64 * 64
+    # Calculate scale factor if image is too large
+    scale_factor = 1.0
+    if max(w, h) > max_size:
+        scale_factor = max_size / max(w, h)
+        new_w = int(w * scale_factor)
+        new_h = int(h * scale_factor)
+        image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+        print(f"Resized from {w}x{h} to {new_w}x{new_h} (scale: {scale_factor:.3f})")
+        w, h = new_w, new_h
     
-    if (new_w, new_h) != (w, h):
-        padded_img = Image.new('RGB', (new_w, new_h), (0, 0, 0))
+    # Pad to multiple of 64
+    pad_w = (w + 63) // 64 * 64
+    pad_h = (h + 63) // 64 * 64
+    
+    if (pad_w, pad_h) != (w, h):
+        padded_img = Image.new('RGB', (pad_w, pad_h), (0, 0, 0))
         padded_img.paste(image, (0, 0))
-        return padded_img, original_size
+        return padded_img, original_size, scale_factor
     
-    return image, original_size
+    return image, original_size, scale_factor
 
 
-def postprocess_image(image: Image.Image, original_size: tuple[int, int]) -> Image.Image:
-    """Crop image back to original size."""
-    if image.size != original_size:
-        return image.crop((0, 0, original_size[0], original_size[1]))
+def postprocess_image(image: Image.Image, original_size: tuple[int, int], scale_factor: float) -> Image.Image:
+    """Crop padding and resize back to original dimensions."""
+    w, h = image.size
+    original_w, original_h = original_size
+    
+    # First crop to the scaled size (remove padding)
+    if scale_factor < 1.0:
+        scaled_w = int(original_w * scale_factor)
+        scaled_h = int(original_h * scale_factor)
+        image = image.crop((0, 0, scaled_w, scaled_h))
+        # Then resize back to original
+        image = image.resize((original_w, original_h), Image.Resampling.LANCZOS)
+        print(f"Upscaled back to original size: {original_w}x{original_h}")
+    else:
+        # Just crop to original size
+        if image.size != original_size:
+            image = image.crop((0, 0, original_w, original_h))
+    
     return image
 
 
@@ -104,12 +133,12 @@ def remove_watermark(
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Preprocess image - pad to multiple of 64
-    processed_image, _ = preprocess_image(input_image)
+    # Preprocess image - resize if too large and pad to multiple of 64
+    processed_image, original_size, scale_factor = preprocess_image(input_image)
     padded_w, padded_h = processed_image.size
-    print(f"Padded image size: {padded_w}x{padded_h}")
+    print(f"Processed image size: {padded_w}x{padded_h} (scale: {scale_factor:.3f})")
     
-    # Run regeneration with explicit dimensions
+    # Run regeneration
     result = pipe(
         prompt="",  # Empty prompt for pure regeneration
         image=processed_image,
@@ -117,16 +146,13 @@ def remove_watermark(
         num_inference_steps=num_inference_steps,
         guidance_scale=0.0,  # No guidance for pure regeneration
         generator=generator,
-        width=padded_w,
-        height=padded_h,
     ).images[0]
     
     print(f"Pipeline output size: {result.size}")
     
-    # Crop back to ORIGINAL size (not padded size)
-    if result.size != (original_w, original_h):
-        result = result.crop((0, 0, original_w, original_h))
-        print(f"Cropped to original size: {result.size}")
+    # Postprocess - crop padding and resize back to original
+    result = postprocess_image(result, original_size, scale_factor)
+    print(f"Final output size: {result.size}")
     
     return result, seed