initial demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Gorluxor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -4,11 +4,33 @@ emoji: 🌖
 colorFrom: yellow
 colorTo: indigo
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 4.44.1 # newest at the time 5.44.1
+suggested_hardware: "a10g-large"
+python_version: 3.10
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 short_description: Fine-Grained Image Editing using Pre-Trained Diffusion Model
+tags:
+  - text-to-image
+  - part-editing
+  - stable-diffusion
+  - diffusion-models
+  - transformers
+  - pytorch
+  - computer-vision
+  - generative-ai
+  - image-generation
+  - image-editing
+  - gradio
+  - machine-learning
+  - deep-learning
+  - ai
+  - artificial-intelligence
+  - demo
+  - research
+preload_from_hub:
+    - stabilityai/stable-diffusion-xl-base-1.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,7 +1,291 @@
+#!/usr/bin/env python
+
+import os
+import random
+from typing import Optional, Tuple, Union, List
+
+import numpy as np
+import PIL.Image
 import gradio as gr
+import torch
+import spaces  # 👈 ZeroGPU support
+
+from model import PartEditSDXLModel, PART_TOKENS
+from datasets import load_dataset
+
+import base64
+from io import BytesIO
+import tempfile
+import uuid
+
+MAX_SEED = np.iinfo(np.int32).max
+CACHE_EXAMPLES = os.environ.get("CACHE_EXAMPLES") == "1"
+AVAILABLE_TOKENS = list(PART_TOKENS.keys())
+
+# Download examples directly from the huggingface PartEdit-Bench
+# Login using e.g. `huggingface-cli login` or `hf login` if needed.
+bench = load_dataset("Aleksandar/PartEdit-Bench", revision="v1.1", split="synth")
+
+use_examples = None  # all with None
+logo = "assets/partedit.png"
+loaded_logo = PIL.Image.open(logo).convert("RGB")
+# base encoded logo 
+
+logo_encoded = None
+with open(logo, "rb") as f:
+    logo_encoded = base64.b64encode(f.read()).decode()
+
+
+def _save_image_for_download(edited: Union[PIL.Image.Image, np.ndarray, str, List]) -> str:
+    """Save the first edited image to a temp file and return its filepath."""
+    # clone to be sure we don't modify the input
+    edited = edited.copy() 
+    img = edited[0] if isinstance(edited, list) else edited
+    if isinstance(img, str):
+        # path on disk already
+        return img
+    if isinstance(img, np.ndarray):
+        img = PIL.Image.fromarray(img)
+    assert isinstance(img, PIL.Image.Image), "Edited output must be PIL, ndarray, str path, or list of these."
+    out_path = os.path.join(tempfile.gettempdir(), f"partedit_{uuid.uuid4().hex}.png")
+    img.save(out_path)
+    return out_path
+
+
+
+def get_example(idx, bench):
+    # [prompt_original, subject, token_cls, edit, "", 50, 7.5, seed, 50]
+    example = bench[idx]
+    return [
+        example["prompt_original"],
+        example["subject"],
+        example["token_cls"],
+        example["edit"],
+        "",
+        50,
+        7.5,
+        example["seed"],
+        50,
+    ]
+
+examples = [get_example(idx, bench) for idx in (use_examples if use_examples is not None else range(len(bench)))]
+first_ex = examples[0] if len(examples) else ["", "", AVAILABLE_TOKENS[0], "", "", 50, 7.5, 0, 50]
+
+title = f"""
+<div style="display: flex; align-items: center;">
+    <img src="data:image/png;base64,{logo_encoded}" alt="PartEdit Logo">
+    <div style="margin-left: 10px;">
+      <h1 style="margin: 0;">PartEdit with SDXL</h1>
+      <p style="margin: 2px 0 0 0;">Official demo for the PartEdit paper.</p>
+      <h2 style="margin: 6px 0 0 0;">PartEdit: Fine-Grained Image Editing using Pre-Trained Diffusion Models</h2>
+      <p style="margin: 6px 0 0 0; font-size: 14px;">
+        It <b>simultaneously predicts the part-localization mask and edits the original trajectory</b>.
+        Supports <b>Hugging Face ZeroGPU</b> and one-click <b>Duplicate</b> for private use.
+      </p>
+    </div>
+</div>
+"""
+
+
+def _as_gallery(edited: Union[PIL.Image.Image, np.ndarray, str, List]) -> List:
+    """Ensure the output fits a Gallery component."""
+    if isinstance(edited, list):
+        return edited
+    return [edited]
+
+
+def edit_demo(model: PartEditSDXLModel) -> gr.Blocks:
+    @spaces.GPU(duration=120)  # 👈 request a ZeroGPU allocation during this call
+    def run(
+        prompt: str,
+        subject: str,
+        part: str,
+        edit: str,
+        negative_prompt: str,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        seed: int = 0,
+        t_e: int = 50,
+        progress=gr.Progress(track_tqdm=True),
+    ) -> Tuple[List, Optional[PIL.Image.Image]]:
+        if seed == -1:
+            seed = random.randint(0, MAX_SEED)
+
+        out = model.edit(
+            prompt=prompt,
+            subject=subject,
+            part=part,
+            edit=edit,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            t_e=t_e,
+        )
+
+        # Accept either (image, mask) or just image from model.edit
+        if isinstance(out, tuple) and len(out) == 2:
+            edited, mask_img = out
+        else:
+            edited, mask_img = out, None
+
+        download_path = _save_image_for_download(edited)
+        return _as_gallery(edited), mask_img, gr.update(value=download_path, visible=True)
+
+    
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Group():
+                    prompt = gr.Textbox(
+                        first_ex[0],  # <- was "a closeup of a man full-body"
+                        placeholder="Prompt",
+                        label="Original Prompt",
+                        show_label=True,
+                        max_lines=1,
+                    )
+                    with gr.Row():
+                        subject = gr.Textbox(value=first_ex[1], label="Subject", show_label=True, max_lines=1)
+                        edit = gr.Textbox(value=first_ex[3], label="Edit", show_label=True, max_lines=1)
+                        part = gr.Dropdown(label="Part Name", choices=AVAILABLE_TOKENS, value=first_ex[2])
+
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=int(first_ex[7]))
+                    run_button = gr.Button("Apply Edit")
+
+                with gr.Accordion("Advanced options", open=False):
+                    negative_prompt = gr.Textbox(label="Negative prompt", value=first_ex[4])
+                    num_inference_steps = gr.Slider(
+                        label="Number of steps",
+                        minimum=1,
+                        maximum=PartEditSDXLModel.MAX_NUM_INFERENCE_STEPS,
+                        step=1,
+                        value=int(first_ex[5]),
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance scale",
+                        minimum=0.1,
+                        maximum=30.0,
+                        step=0.1,
+                        value=float(first_ex[6]),
+                    )
+                    t_e = gr.Slider(
+                        label="Editing steps",
+                        minimum=1,
+                        maximum=PartEditSDXLModel.MAX_NUM_INFERENCE_STEPS,
+                        step=1,
+                        value=int(first_ex[8]),
+                    )
+                with gr.Accordion('Citation', open=True):
+                    gr.Markdown(citation)
+
+
+            with gr.Column(scale=3):
+                with gr.Row(equal_height=False):
+                    with gr.Column(scale=1, min_width=120):
+                        mask = gr.Image(label="Editing Mask", width=100, height=100, show_label=True)
+                    with gr.Column(scale=7):
+                        result = gr.Gallery(
+                            label="Edited Image",
+                            height=700,
+                            object_fit="fill",
+                            preview=True,
+                            selected_index=0,
+                            show_label=True,
+                        )
+                download_btn = gr.File(
+                    label="Download full-resolution",
+                    type="filepath",
+                    file_count="single",   # <-- keeps it to one file
+                    interactive=False,
+                    height=48,             # <-- compact
+                    visible=False          # <-- hide until we have a file
+                )
+
+        inputs = [prompt, subject, part, edit, negative_prompt, num_inference_steps, guidance_scale, seed, t_e]
+
+        gr.Examples(
+            examples=examples,
+            inputs=inputs,
+            outputs=[result, mask, download_btn],
+            fn=run,
+            cache_examples=CACHE_EXAMPLES,
+        )
+
+        run_button.click(fn=run, inputs=inputs, outputs=[result, mask, download_btn], api_name="run")
+
+    return demo
+
+
+badges_text = r"""
+<div style="text-align: center; display: flex; justify-content: center; gap: 5px; flex-wrap: wrap;">
+  <a href="https://gorluxor.github.io/part-edit/">
+    <img alt="Project Page" src="https://img.shields.io/badge/%F0%9F%8C%90%20Project%20Page-PartEdit-blue">
+  </a>
+  <a href="https://arxiv.org/abs/2502.04050">
+    <img alt="arXiv" src="https://img.shields.io/badge/arXiv-2502.04050-b31b1b.svg">
+  </a>
+  <a href="https://huggingface.co/datasets/Aleksandar/PartEdit-Bench">
+    <img alt="HF Dataset: PartEdit-Bench" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-PartEdit--Bench-blue">
+  </a>
+  <a href="https://huggingface.co/datasets/Aleksandar/PartEdit-extra">
+    <img alt="HF Dataset: PartEdit-extra" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-PartEdit--extra-blue">
+  </a>
+  <a href="https://s2025.siggraph.org/">
+    <img alt="SIGGRAPH 2025" src="https://img.shields.io/badge/%F0%9F%8E%A8%20Accepted-SIGGRAPH%202025-blueviolet">
+  </a>
+  <a href="https://github.com/Gorluxor/partedit/blob/main/LICENSE">
+    <img alt="Code License" src="https://img.shields.io/badge/license-MIT-blue.svg">
+  </a>
+</div>
+""".strip()
+
+citation = r"""
+If you use this demo, please cite the following paper:
+```
+@inproceedings{cvejic2025partedit,
+  title={PartEdit: Fine-Grained Image Editing using Pre-Trained Diffusion Models},
+  author={Cvejic, Aleksandar and Eldesokey, Abdelrahman and Wonka, Peter},
+  booktitle={Proceedings of the Special Interest Group on Computer Graphics and Interactive Techniques Conference Conference Papers},
+  pages={1--11},
+  year={2025}
+}
+```
+"""
+
+DESCRIPTION = title + badges_text
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. On ZeroGPU Spaces, a GPU will be requested when you click <b>Apply Edit</b>.</p>"
+
+def running_in_hf_space() -> bool:
+    # Common env vars present on Hugging Face Spaces
+    return (
+        os.getenv("SYSTEM") == "spaces" or
+        any(os.getenv(k) for k in (
+            "SPACE_ID", "HF_SPACE_ID", "SPACE_REPO_ID",
+            "SPACE_REPO_NAME", "SPACE_AUTHOR_NAME", "SPACE_TITLE"
+        ))
+    )
+
+if __name__ == "__main__":
+    model = PartEditSDXLModel()
+
+    with gr.Blocks(css="style.css") as demo:
+        gr.Markdown(DESCRIPTION)
+
+        # Always show Duplicate button on Spaces
+        gr.DuplicateButton(
+            value="Duplicate Space for private use",
+            elem_id="duplicate-button",
+            variant="huggingface",
+            size="lg",
+            visible=running_in_hf_space(),
+        )
 
-def greet(name):
-    return "Hello " + name + "!!"
+        # Single tab: PartEdit only
+        with gr.Tabs():
+            with gr.Tab(label="PartEdit", id="edit"):
+                edit_demo(model)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+        demo.queue(max_size=20).launch()

model.py

@@ -0,0 +1,136 @@
+import gc
+
+import PIL.Image
+import torch
+
+from stable_diffusion_xl_partedit import PartEditPipeline, DotDictExtra, Binarization, PaddingStrategy, EmptyControl
+from diffusers import AutoencoderKL
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import CLIPImageProcessor
+
+from huggingface_hub import hf_hub_download
+
+available_pts = [
+    "pt/torso_custom.pt", # this is human torso only
+    "pt/chair_custom.pt", # this is seat of the chair only
+    "pt/carhood_custom.pt", 
+    "pt/partimage_biped_head.pt", # this is essentially monkeys
+    "pt/partimage_carbody.pt", # this is everything except the wheels
+    "pt/partimage_human_hair.pt", 
+    "pt/partimage_human_head.pt", # this is essentially faces
+    "pt/partimage_human_torso.pt", # use custom on in favour of this one
+    "pt/partimage_quadruped_head.pt", # this is essentially animals on 4 legs
+]
+
+def download_part(index):
+    return hf_hub_download(
+        repo_id="Aleksandar/PartEdit-extra",
+        repo_type="dataset",
+        filename=available_pts[index]
+    )
+
+PART_TOKENS = {
+    "human_head": download_part(6),
+    "human_hair": download_part(5),
+    "human_torso_custom": download_part(0), # custom one
+    "chair_custom": download_part(1),
+    "carhood_custom": download_part(2),
+    "carbody": download_part(4),
+    "biped_head": download_part(8),
+    "quadruped_head": download_part(3),
+    "human_torso": download_part(7), # based on partimage
+}
+
+
+class PartEditSDXLModel:
+    MAX_NUM_INFERENCE_STEPS = 50
+
+    def __init__(self):
+        if torch.cuda.is_available():
+            self.device = torch.device(f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu")
+            self.sd_pipe, self.partedit_pipe = PartEditPipeline.default_pipeline(self.device)
+        else:
+            self.pipe = None
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        seed: int = 0,
+        eta: float = 0,
+    ) -> PIL.Image.Image:
+
+        if not torch.cuda.is_available():
+            raise RuntimeError("This demo does not work on CPU!")
+
+        out = self.sd_pipe(
+            prompt=prompt,
+            # negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            eta=eta,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        return out
+
+    def edit(
+        self,
+        prompt: str,
+        subject: str,
+        part: str,
+        edit: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        seed: int = 0,
+        eta: int = 0,
+        t_e: int = 50,
+    ) -> PIL.Image.Image:
+
+        # Sanity Checks
+        if not torch.cuda.is_available():
+            raise RuntimeError("This demo does not work on CPU!")
+
+        if part in PART_TOKENS:
+            token_path = PART_TOKENS[part]
+        else:
+            raise ValueError(f"Part `{part}` is not supported!")
+
+        if subject not in prompt:
+            raise ValueError(f"The subject `{subject}` does not exist in the original prompt!")
+
+        prompts = [
+            prompt,
+            prompt.replace(subject, edit),
+        ]
+
+        # PartEdit Parameters
+        cross_attention_kwargs = {
+            "edit_type": "replace",
+            "n_self_replace": 0.0,
+            "n_cross_replace": {"default_": 1.0, edit: 0.4},
+        }
+        extra_params = DotDictExtra()
+        extra_params.update({"omega": 1.5, "edit_steps": t_e})
+
+        out = self.partedit_pipe(
+            prompt=prompts,
+            # negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            eta=eta,
+            generator=torch.Generator().manual_seed(seed),
+            cross_attention_kwargs=cross_attention_kwargs,
+            extra_kwargs=extra_params,
+            embedding_opt=token_path,
+        ).images[:2][::-1]
+
+        mask = self.partedit_pipe.visualize_map_across_time()
+        gc.collect()
+        torch.cuda.empty_cache()
+        return out, mask

pyproject.toml

@@ -0,0 +1,5 @@
+[tool.ruff]
+extend-select = ["C4", "SIM", "TCH"]
+ignore = ["F401"]
+show-fixes = true
+target-version = "py39"

requirements.txt

@@ -0,0 +1,31 @@
+# Requires Python 3.10.x
+# Pull PyTorch CUDA 11.8 wheels
+--extra-index-url https://download.pytorch.org/whl/cu118
+
+setuptools>=61.0
+numpy<1.24
+# ipywidgets
+# black[jupyter]
+# jupyterlab
+# matplotlib
+einops
+ftfy
+regex
+tqdm
+
+# Core ML stack (PyTorch 2.1.0 + CUDA 11.8)
+torch==2.1.0
+torchvision==0.16.0
+# torchaudio==2.1.0
+
+# UI / HF stack
+gradio<5.0           # tested on 4.29; should work on 4.44.1 with pydantic fix
+huggingface_hub<0.26.0
+pydantic<=2.10.6
+
+# Diffusion / training utils
+diffusers==0.27.2
+transformers==4.44.1
+accelerate
+datasets             # tested on 3.3.2
+spaces