app.py

"""
SD3.5 Watermark Remover - Hugging Face Space

Uses Stable Diffusion 3.5 img2img regeneration to remove watermarks
while preserving image quality. Powered by ZeroGPU.

Based on: https://github.com/XuandongZhao/WatermarkAttacker
"""

import os
import gradio as gr
import numpy as np
import spaces
import torch
from PIL import Image
from huggingface_hub import login

# Login with HF_TOKEN for gated model access
if os.environ.get("HF_TOKEN"):
    login(token=os.environ["HF_TOKEN"])
    print("Logged in with HF_TOKEN")

from diffusers import StableDiffusion3Img2ImgPipeline

# Model configuration
MODEL_ID = "stabilityai/stable-diffusion-3.5-large"
MAX_IMAGE_SIZE = 1536
MAX_SEED = np.iinfo(np.int32).max

# Determine device and dtype
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

# Load model at startup (will use ZeroGPU when decorated function is called)
print(f"Loading SD3.5 model: {MODEL_ID}")
pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
)
pipe = pipe.to(device)
print("Model loaded successfully!")


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
    
    # 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
    
    # 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, scale_factor
    
    return image, original_size, scale_factor


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


@spaces.GPU(duration=90)
def remove_watermark(
    input_image: Image.Image,
    strength: float = 0.3,
    num_inference_steps: int = 28,
    seed: int = 42,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[Image.Image, int]:
    """
    Remove watermark from image using SD3.5 img2img regeneration.
    
    Args:
        input_image: Input image with watermark
        strength: Denoising strength (0.0-1.0). Lower = higher quality, less change.
        num_inference_steps: Number of denoising steps
        seed: Random seed for reproducibility
        randomize_seed: Whether to randomize the seed
        
    Returns:
        Tuple of (output image, seed used)
    """
    if input_image is None:
        raise gr.Error("Please upload an image first!")
    
    # Convert to RGB if needed
    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 - 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"Processed image size: {padded_w}x{padded_h} (scale: {scale_factor:.3f})")
    
    # Run regeneration
    result = pipe(
        prompt="",  # Empty prompt for pure regeneration
        image=processed_image,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=0.0,  # No guidance for pure regeneration
        generator=generator,
    ).images[0]
    
    print(f"Pipeline output 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


# Example images (you can add your own)
examples = [
    ["examples/watermarked_sample.png", 0.3, 28, 42, True],
]

# Custom CSS
css = """
#col-container {
    margin: 0 auto;
    max-width: 900px;
}

.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}

footer {
    visibility: hidden;
}
"""

# Build Gradio interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # 🎨 SD3.5 Watermark Remover
        
        Remove watermarks from images using **Stable Diffusion 3.5** regeneration.
        
        This tool uses img2img to regenerate the image while preserving its content,
        effectively removing watermarks without manual editing.
        
        **How it works:** Lower strength values produce higher quality outputs closer to the original.
        Start with 0.3 and adjust as needed.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Input Image (with watermark)",
                    type="pil",
                    height=400,
                )
                
                run_button = gr.Button("🚀 Remove Watermark", variant="primary", size="lg")
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    strength = gr.Slider(
                        label="Strength",
                        info="Lower = higher quality, less change. Higher = more aggressive removal.",
                        minimum=0.1,
                        maximum=0.6,
                        step=0.05,
                        value=0.3,
                    )
                    
                    num_inference_steps = gr.Slider(
                        label="Inference Steps",
                        info="More steps = higher quality, slower processing.",
                        minimum=10,
                        maximum=50,
                        step=1,
                        value=28,
                    )
                    
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    
                    randomize_seed = gr.Checkbox(
                        label="Randomize seed",
                        value=True,
                    )
            
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Output Image (watermark removed)",
                    type="pil",
                    height=400,
                )
                
                output_seed = gr.Number(label="Seed Used", interactive=False)
        
        gr.Markdown("""
        ### 💡 Tips
        - **Strength 0.2-0.3**: Best for subtle watermarks, preserves most detail
        - **Strength 0.4-0.5**: Better for prominent watermarks, may alter some details
        - **Increase steps**: If quality is poor, try 35-40 steps
        
        ---
        
        Based on [WatermarkAttacker](https://github.com/XuandongZhao/WatermarkAttacker) research.
        Uses Stable Diffusion 3.5 Large from [Stability AI](https://huggingface.co/stabilityai/stable-diffusion-3.5-large).
        """)
    
    # Connect events
    gr.on(
        triggers=[run_button.click],
        fn=remove_watermark,
        inputs=[
            input_image,
            strength,
            num_inference_steps,
            seed,
            randomize_seed,
        ],
        outputs=[output_image, output_seed],
    )


if __name__ == "__main__":
    demo.launch()