| import spaces |
| import gradio as gr |
| import pyte |
| import math |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| import os |
| import pty |
| import select |
| import fcntl |
| import struct |
| import termios |
| from pathlib import Path |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
| from ultralytics import YOLO |
| from preservation import PreservedWatermarkRemover |
| from ignored_preview import save_ignored_preview |
| from clean import load_clean, clean_image |
| from illustrious import load_illustrious, prepare_illustrious, infer_illustrious, save_illustrious |
| from watermark_remover import WatermarkRemover |
|
|
| INPUT_IMAGE_HEIGHT = 600 |
| TEXT_WORKER = [sys.executable, "-m", "text_detection"] |
| TIME_GPU = 13 |
| MODEL_DEVICE = "cuda" |
| PART_MODELS = { |
| "eyes": ("killjoyelite/anime-eye-yolov8", "eye_yolov8n.pt", 0.25, 0), |
| "hands": ("deepghs/anime_hand_detection", "hand_detect_v1.0_s/model.pt", 0.3, 0.12), |
| "feet": ("Claquasse/foot_anime_yolo", "foot_anime_yolo11m_v3.pt", 0.25, 0.02), |
| } |
|
|
|
|
| def load_models(): |
| subprocess.run(TEXT_WORKER, check=True) |
| removers = {} |
| for model_id in (WatermarkRemover.DEFAULT_MODEL_ID, WatermarkRemover.CTRLREGEN_MODEL_ID): |
| remover = PreservedWatermarkRemover((), model_id=model_id, device=MODEL_DEVICE) |
| remover.preload() |
| removers[model_id] = remover |
| parts = {} |
| for part, (repo, filename, _, _) in PART_MODELS.items(): |
| detector = YOLO(hf_hub_download(repo_id=repo, filename=filename)) |
| detector.fuse() |
| parts[part] = detector.to(MODEL_DEVICE) |
| return { |
| "watermark": removers, |
| "parts": parts, |
| "clean": load_clean(MODEL_DEVICE), |
| "illustrious": load_illustrious(MODEL_DEVICE), |
| } |
|
|
|
|
| def set_terminal_size(fd, rows=24, cols=80): |
| size = struct.pack('HHHH', rows, cols, 0, 0) |
| fcntl.ioctl(fd, termios.TIOCSWINSZ, size) |
|
|
|
|
| def run_in_pty(cmd, env=None, timeout=600): |
| screen = pyte.Screen(80, 30) |
| stream = pyte.Stream(screen) |
|
|
| master_fd, slave_fd = pty.openpty() |
| set_terminal_size(master_fd, rows=30, cols=80) |
|
|
| proc = subprocess.Popen( |
| cmd, stdout=slave_fd, stderr=slave_fd, stdin=slave_fd, |
| close_fds=True, env=env, |
| ) |
| os.close(slave_fd) |
|
|
| start_time = time.time() |
| try: |
| while True: |
| if time.time() - start_time > timeout: |
| proc.kill() |
| yield get_screen_text(screen), True |
| return |
|
|
| ready, _, _ = select.select([master_fd], [], [], 0.3) |
| if ready: |
| try: |
| data = os.read(master_fd, 4096) |
| if not data: |
| break |
| stream.feed(data.decode("utf-8", errors="replace")) |
| except OSError: |
| break |
|
|
| if proc.poll() is not None: |
| while True: |
| r, _, _ = select.select([master_fd], [], [], 0.1) |
| if r: |
| try: |
| d = os.read(master_fd, 4096) |
| if not d: |
| break |
| stream.feed(d.decode("utf-8", errors="replace")) |
| except OSError: |
| break |
| else: |
| break |
| yield get_screen_text(screen), True |
| return |
|
|
| yield get_screen_text(screen), False |
| |
| |
| |
| yield get_screen_text(screen), True |
| finally: |
| try: |
| os.close(master_fd) |
| except OSError: |
| pass |
| if proc.poll() is None: |
| proc.kill() |
| proc.wait() |
|
|
|
|
| def get_screen_text(screen): |
| lines = [] |
| for row in range(screen.lines): |
| line = "" |
| for col in range(screen.columns): |
| line += screen.buffer[row][col].data |
| lines.append(line.rstrip()) |
| while lines and not lines[-1].strip(): |
| lines.pop() |
| return '\n'.join(lines) |
|
|
|
|
| def escape_html(text): |
| return text.replace('&', '&').replace('<', '<').replace('>', '>') |
|
|
|
|
| def make_terminal_html(content, cursor=True): |
| cursor_html = "<span class='cursor'>█</span>" if cursor else "" |
| return f"<div class='terminal'><pre>{content}{cursor_html}</pre></div>" |
|
|
|
|
| @spaces.GPU(duration=TIME_GPU) |
| def run_watermark(input_path, output_path, strength, steps, model_id, boxes, text_mask): |
| selected = MODELS["watermark"][model_id or WatermarkRemover.DEFAULT_MODEL_ID] |
| selected.boxes = boxes |
| selected.text_mask = text_mask |
| try: |
| selected.remove_watermark(input_path, output_path, strength=strength, num_inference_steps=steps) |
| finally: |
| selected.boxes = () |
| selected.text_mask = None |
|
|
|
|
| @spaces.GPU(duration=TIME_GPU) |
| def run_illustrious(image, control, mask, strength, steps): |
| return infer_illustrious(MODELS["illustrious"], image, control, mask, strength, steps) |
|
|
|
|
| @spaces.GPU(duration=TIME_GPU) |
| def run_clean(input_path, output_path): |
| clean_image(MODELS["clean"], input_path, output_path) |
|
|
|
|
| @spaces.GPU(duration=TIME_GPU) |
| def detect_parts(input_path, part): |
| _, _, confidence, padding = PART_MODELS[part] |
| model = MODELS["parts"][part] |
| result = model.predict(source=str(input_path), conf=confidence, device=MODEL_DEVICE, verbose=False)[0] |
| width, height = result.orig_shape[1], result.orig_shape[0] |
| boxes = [] |
| for x1, y1, x2, y2 in result.boxes.xyxy.cpu().tolist(): |
| dx, dy = (x2 - x1) * padding, (y2 - y1) * padding |
| box = ( |
| max(0, math.floor(x1 - dx)), max(0, math.floor(y1 - dy)), |
| min(width, math.ceil(x2 + dx)), min(height, math.ceil(y2 + dy)), |
| ) |
| if box[2] > box[0] and box[3] > box[1]: |
| boxes.append(box) |
| return boxes |
|
|
|
|
| def detect_text(input_path): |
| with tempfile.TemporaryDirectory() as directory: |
| output_path = Path(directory) / "mask.png" |
| subprocess.run([*TEXT_WORKER, str(input_path), str(output_path)], check=True) |
| with Image.open(output_path) as mask: |
| return mask.copy() |
|
|
|
|
| def restore_parts(input_path, output_path, boxes, text_mask=None): |
| with Image.open(input_path) as original, Image.open(output_path) as cleaned: |
| result = cleaned.copy() |
| width, height = result.size |
| sx, sy = width / original.width, height / original.height |
| original = original.convert(result.mode) |
| if original.size != result.size: |
| original = original.resize(result.size, Image.Resampling.LANCZOS) |
| for x1, y1, x2, y2 in boxes: |
| box = ( |
| max(0, math.floor(x1 * sx)), max(0, math.floor(y1 * sy)), |
| min(width, math.ceil(x2 * sx)), min(height, math.ceil(y2 * sy)), |
| ) |
| if box[2] > box[0] and box[3] > box[1]: |
| result.paste(original.crop(box), box) |
| if text_mask is not None: |
| result.paste(original, (0, 0), text_mask.resize(result.size, Image.Resampling.NEAREST)) |
| result.save(output_path) |
|
|
|
|
| def remove_watermark(file_obj, strength, steps, model, ignore_eyes=True, ignore_hands=True, |
| ignore_feet=True, clean_after=True, ignore_text=True): |
| if file_obj is None: |
| raise gr.Error("Please upload an image first.") |
|
|
| steps = max(int(steps), math.ceil(1 / strength)) |
| input_path = Path(file_obj) |
| tmp_dir = tempfile.mkdtemp() |
| output_path = Path(tmp_dir) / "cleaned.png" |
| preview_path = Path(tmp_dir) / "ignored.png" |
| work_path = Path(tmp_dir) / "result.png" |
| enabled = {"eyes": ignore_eyes, "hands": ignore_hands, "feet": ignore_feet} |
| regions = {} |
| text_mask = None |
|
|
| if model == "CtrlRegen Illustrious (SDXL)": |
| cmd_display = f"CtrlRegen-like Illustrious: {input_path.name} --strength {strength} --steps {int(steps)} -o cleaned.png" |
| else: |
| profile = " --model-profile ctrlregen" if model == "CtrlRegen" else "" |
| cmd_display = f"$ noai-watermark {input_path.name}{profile} --strength {strength} --steps {int(steps)} -o cleaned.png" |
| cmd_html = f"<span style='color:#8be9fd'>{escape_html(cmd_display)}</span>" |
|
|
| try: |
| for part, ignore in enabled.items(): |
| if ignore: |
| regions[part] = detect_parts(input_path, part) |
| if ignore_text: |
| text_mask = detect_text(input_path) |
| if text_mask.getbbox() is None: |
| text_mask = None |
| boxes = [box for found in regions.values() for box in found] |
| if model == "CtrlRegen Illustrious (SDXL)": |
| image, control, mask, crop, size = prepare_illustrious(input_path, boxes, text_mask) |
| result = run_illustrious(image, control, mask, strength, int(steps)) |
| save_illustrious(result, crop, size, work_path) |
| else: |
| model_id = WatermarkRemover.CTRLREGEN_MODEL_ID if model == "CtrlRegen" else None |
| run_watermark(input_path, work_path, strength, int(steps), model_id, boxes, text_mask) |
| with Image.open(work_path) as image: |
| image.save(output_path) |
| if clean_after: |
| run_clean(output_path, output_path) |
| if boxes or text_mask is not None: |
| restore_parts(input_path, output_path, boxes, text_mask) |
| save_ignored_preview(input_path, preview_path, regions, text_mask) |
| except Exception as error: |
| return make_terminal_html(cmd_html + "\n\n" + escape_html(str(error)), cursor=False), None, None |
|
|
| preserved = "".join( |
| f" Preserved {len(boxes)} detected {part}." if boxes else f" No {part} detected; none preserved." |
| for part, boxes in regions.items() |
| ) |
| if ignore_text: |
| preserved += " Preserved text character mask." if text_mask is not None else " No text characters segmented; none preserved." |
| cleaned = " Applied Clean after (1x Archivist Soft)." if clean_after else "" |
| return make_terminal_html(cmd_html + f"\n\nDone!{preserved}{cleaned}", cursor=False), str(output_path), str(preview_path) |
|
|
|
|
| def check_ai_metadata(file_obj): |
| if file_obj is None: |
| raise gr.Error("Please upload an image first.") |
|
|
| input_path = Path(file_obj) |
|
|
| env = os.environ.copy() |
| env["PYTHONUNBUFFERED"] = "1" |
| env["TERM"] = "xterm-256color" |
| env["FORCE_COLOR"] = "1" |
| env["CLICOLOR_FORCE"] = "1" |
|
|
| final_text = "" |
| for screen_text, is_done in run_in_pty( |
| ["noai-watermark", str(input_path), "--check-ai"], |
| env=env, timeout=30 |
| ): |
| final_text = screen_text |
|
|
| display = escape_html(final_text) |
| cmd_display = f"$ noai-watermark {input_path.name} --check-ai" |
| return make_terminal_html(f"<span style='color:#8be9fd'>{escape_html(cmd_display)}</span>\n\n" + display, cursor=False) |
|
|
|
|
| |
| terminal_css = """ |
| .terminal { |
| background: #1a1b26; |
| color: #a9b1d6; |
| border-radius: 10px; |
| padding: 20px 24px; |
| font-family: 'JetBrains Mono', 'Fira Code', 'SF Mono', 'Cascadia Code', 'Consolas', monospace; |
| font-size: 13px; |
| line-height: 1.7; |
| overflow-x: auto; |
| border: 1px solid #33354a; |
| max-height: 500px; |
| overflow-y: auto; |
| box-shadow: 0 4px 16px rgba(0,0,0,0.5); |
| } |
| .terminal pre { |
| margin: 0; |
| white-space: pre-wrap; |
| word-wrap: break-word; |
| color: #a9b1d6; |
| font-family: inherit; |
| } |
| .terminal .cursor { |
| color: #7aa2f7; |
| animation: blink 1s step-end infinite; |
| } |
| @keyframes blink { |
| 50% { opacity: 0; } |
| } |
| """ |
|
|
| MODELS = load_models() |
|
|
| with gr.Blocks(title="noai-watermark") as demo: |
|
|
| gr.HTML(""" |
| <h1 style="text-align:center;">🛡️ noai-watermark</h1> |
| <p style="text-align:center; color:#888;"> |
| Remove invisible AI watermarks (SynthID, StableSignature, TreeRing) |
| and strip AI metadata from images |
| </p> |
| """) |
|
|
| with gr.Tabs(): |
|
|
| with gr.TabItem("🔧 Watermark Removal"): |
| upload_wm = gr.Image(label="Upload Image (PNG or JPEG)", sources=["upload"], |
| type="filepath", image_mode=None, format="png", |
| height=INPUT_IMAGE_HEIGHT) |
| model = gr.Radio( |
| ["CtrlRegen Illustrious (SDXL)", "CtrlRegen", "Default"], |
| value="CtrlRegen", label="Model", |
| ) |
| with gr.Row(): |
| ignore_eyes = gr.Checkbox(value=True, label="Ignore eyes") |
| ignore_hands = gr.Checkbox(value=True, label="Ignore hands") |
| ignore_feet = gr.Checkbox(value=True, label="Ignore feet") |
| ignore_text = gr.Checkbox(value=True, label="Ignore text") |
| clean_after = gr.Checkbox(value=True, label="Clean after") |
|
|
| with gr.Row(): |
| strength = gr.Slider( |
| minimum=0.01, maximum=1.0, value=0.02, step=0.01, |
| label="Strength", |
| info="Lower = less change. Increase if watermark persists.", |
| ) |
| steps = gr.Slider( |
| minimum=10, maximum=100, value=50, step=5, |
| label="Denoising Steps", |
| info="More steps = better quality but slower.", |
| ) |
|
|
| with gr.Row(): |
| btn_super_minimal = gr.Button("Super Minimal (0.02/50)", size="sm") |
| btn_minimal = gr.Button("Minimal (0.04/50)", size="sm") |
| btn_balanced = gr.Button("Balanced (0.15/50)", size="sm") |
| btn_aggressive = gr.Button("Aggressive (0.35/60)", size="sm") |
| btn_max = gr.Button("Maximum (0.70/60)", size="sm") |
|
|
| btn_super_minimal.click(lambda: (0.02, 50), outputs=[strength, steps]) |
| btn_minimal.click(lambda: (0.04, 50), outputs=[strength, steps]) |
| btn_balanced.click(lambda: (0.15, 50), outputs=[strength, steps]) |
| btn_aggressive.click(lambda: (0.35, 60), outputs=[strength, steps]) |
| btn_max.click(lambda: (0.70, 60), outputs=[strength, steps]) |
|
|
| run_btn = gr.Button("🚀 Remove Watermark", variant="primary", size="lg") |
| wm_terminal = gr.HTML(value=make_terminal_html("<span style='color:#565f89'>Waiting for input...</span>")) |
|
|
| with gr.Row(): |
| output_image = gr.Image( |
| type="filepath", image_mode=None, format="png", interactive=False, |
| show_label=False, height=INPUT_IMAGE_HEIGHT, buttons=["download", "fullscreen"], |
| ) |
| ignored_image = gr.Image( |
| type="filepath", image_mode=None, format="png", interactive=False, |
| show_label=False, height=INPUT_IMAGE_HEIGHT, buttons=["download", "fullscreen"], |
| ) |
|
|
| run_btn.click( |
| fn=remove_watermark, |
| inputs=[upload_wm, strength, steps, model, ignore_eyes, ignore_hands, ignore_feet, clean_after, ignore_text], |
| outputs=[wm_terminal, output_image, ignored_image], |
| concurrency_limit=1, |
| ) |
|
|
| with gr.TabItem("🔍 Check AI Metadata"): |
| upload_meta = gr.Image(label="Upload Image to Check (PNG or JPEG)", sources=["upload"], |
| type="filepath", image_mode=None, |
| height=INPUT_IMAGE_HEIGHT) |
| check_btn = gr.Button("🔍 Check Metadata", variant="secondary", size="lg") |
| meta_terminal = gr.HTML(value=make_terminal_html("<span style='color:#565f89'>Waiting for input...</span>")) |
|
|
| check_btn.click( |
| fn=check_ai_metadata, |
| inputs=[upload_meta], |
| outputs=[meta_terminal], |
| ) |
|
|
| gr.Markdown(""" |
| --- |
| <div style="text-align:center; color:#888; font-size:0.85em;"> |
| <a href="https://github.com/mertizci/noai-watermark" target="_blank">GitHub</a> · |
| <a href="https://pypi.org/project/noai-watermark/" target="_blank">PyPI</a> · |
| Built by <a href="https://github.com/mertizci" target="_blank">Mert IZCI</a> |
| <br>For security research and defensive analysis only. |
| </div> |
| """) |
|
|
| if __name__ == "__main__": |
| demo.launch(ssr_mode=False, css=terminal_css, allowed_paths=["/tmp"]) |
|
|
