""" 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()