1import os
2import gradio as gr
3import numpy as np
4import spaces
5import torch
6import random
7from PIL import Image
8from typing import Iterable
9
10from diffusers import FluxKontextPipeline
11from diffusers.utils import load_image
12from huggingface_hub import hf_hub_download
13from gradio_imageslider import ImageSlider
14
15device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16
17# --- Main Model Initialization ---
18MAX_SEED = np.iinfo(np.int32).max
19pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
20
21# --- Load New Adapter ---
22pipe.load_lora_weights("prithivMLmods/Kontext-Watermark-Remover", weight_name="Kontext-Watermark-Remover.safetensors", adapter_name="watermark_remover")
23
24
25@spaces.GPU
26def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
27 """
28 Perform image editing, returning a pair for the ImageSlider.
29 """
30 if not input_image:
31 raise gr.Error("Please upload an image for editing.")
32
33 pipe.set_adapters(["watermark_remover"], adapter_weights=[1.0])
34
35 if randomize_seed:
36 seed = random.randint(0, MAX_SEED)
37
38 original_image = input_image.copy().convert("RGB")
39
40 image = pipe(
41 image=original_image,
42 prompt=prompt,
43 guidance_scale=guidance_scale,
44 width = original_image.size[0],
45 height = original_image.size[1],
46 num_inference_steps=steps,
47 generator=torch.Generator().manual_seed(seed),
48 ).images[0]
49
50 return (original_image, image), seed, gr.Button(visible=True)
51
52css="""
53#col-container {
54 margin: 0 auto;
55 max-width: 960px;
56}
57#main-title h1 {font-size: 2.1em !important;}
58"""
59
60with gr.Blocks(css=css) as demo:
61
62 with gr.Column(elem_id="col-container"):
63 gr.Markdown("# **Photo-Mate-i2i: Watermark Remover**", elem_id="main-title")
64 gr.Markdown("Image manipulation with FLUX.1 Kontext. This demo focuses on watermark removal.")
65
66 with gr.Row():
67 with gr.Column():
68 input_image = gr.Image(label="Upload Image with Watermark", type="pil", height="300")
69 with gr.Row():
70 prompt = gr.Text(
71 label="Edit Prompt",
72 show_label=False,
73 max_lines=1,
74 placeholder="e.g., 'Remove the watermark'",
75 container=False,
76 value="[photo content], remove any watermark text or logos from the image while preserving the background, texture, lighting, and overall realism. Ensure the edited areas blend seamlessly with surrounding details, leaving no visible traces of watermark removal."
77 )
78 run_button = gr.Button("Run", variant="primary", scale=0)
79 with gr.Accordion("Advanced Settings", open=False):
80
81 seed = gr.Slider(
82 label="Seed",
83 minimum=0,
84 maximum=MAX_SEED,
85 step=1,
86 value=0,
87 )
88
89 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
90
91 guidance_scale = gr.Slider(
92 label="Guidance Scale",
93 minimum=1,
94 maximum=10,
95 step=0.1,
96 value=2.5,
97 )
98
99 steps = gr.Slider(
100 label="Steps",
101 minimum=1,
102 maximum=30,
103 value=28,
104 step=1
105 )
106
107 with gr.Column():
108 output_slider = ImageSlider(label="Before / After", show_label=False, interactive=False)
109 reuse_button = gr.Button("Reuse this image", visible=False)
110
111 gr.on(
112 triggers=[run_button.click, prompt.submit],
113 fn=infer,
114 inputs=[input_image, prompt, seed, randomize_seed, guidance_scale, steps],
115 outputs=[output_slider, seed, reuse_button]
116 )
117
118 reuse_button.click(
119 fn=lambda images: images[1] if isinstance(images, (list, tuple)) and len(images) > 1 else images,
120 inputs=[output_slider],
121 outputs=[input_image]
122 )
123
124demo.launch(mcp_server=True, ssr_mode=False, show_error=True)