Upload folder using huggingface_hub

app.py

@@ -94,16 +94,22 @@ def remove_watermark(
     if input_image.mode != 'RGB':
         input_image = input_image.convert('RGB')
     
+    # Store original size BEFORE any processing
+    original_w, original_h = input_image.size
+    print(f"Original image size: {original_w}x{original_h}")
+    
     # Handle seed
     if randomize_seed:
         seed = np.random.randint(0, MAX_SEED)
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Preprocess image
-    processed_image, original_size = preprocess_image(input_image)
+    # Preprocess image - pad to multiple of 64
+    processed_image, _ = preprocess_image(input_image)
+    padded_w, padded_h = processed_image.size
+    print(f"Padded image size: {padded_w}x{padded_h}")
     
-    # Run regeneration
+    # Run regeneration with explicit dimensions
     result = pipe(
         prompt="",  # Empty prompt for pure regeneration
         image=processed_image,
@@ -111,10 +117,16 @@ 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]
     
-    # Postprocess to original size
-    result = postprocess_image(result, original_size)
+    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}")
     
     return result, seed