fixed cache path

.gitignore

@@ -0,0 +1,55 @@
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+test/*
+# C extensions
+*.so
+output
+*labels.json
+thirdparty
+Inversion_SDXL
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# LSP config files
+pyrightconfig.json

README.md

@@ -31,6 +31,8 @@ tags:
   - research
 preload_from_hub:
     - stabilityai/stable-diffusion-xl-base-1.0
+    - Aleksandar/PartEdit-Bench
+    - Aleksandar/PartEdit-extra
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -18,13 +18,23 @@ from io import BytesIO
 import tempfile
 import uuid
 
+import pathlib
+HF_BASE = "/data/.huggingface" if os.getenv("SYSTEM") == "spaces" else "./.hf_cache" # /data/hf_cache
+HF_BASE = str(pathlib.Path(HF_BASE).absolute())
+os.environ.setdefault("HF_HOME", HF_BASE)
+os.environ.setdefault("HF_HUB_CACHE", os.path.join(HF_BASE, "hub"))
+os.environ.setdefault("HF_DATASETS_CACHE", os.path.join(HF_BASE, "datasets"))
+os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(HF_BASE, "transformers"))
+for k in ("HF_HUB_CACHE", "HF_DATASETS_CACHE", "TRANSFORMERS_CACHE"):
+    os.makedirs(os.environ[k], exist_ok=True)
+
 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")
+bench = load_dataset("Aleksandar/PartEdit-Bench", revision="v1.1", split="synth", cache_dir=os.environ["HF_DATASETS_CACHE"])
 
 use_examples = None  # all with None
 logo = "assets/partedit.png"
@@ -37,18 +47,14 @@ with open(logo, "rb") as f:
 
 
 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."
+    item = edited[0] if isinstance(edited, list) else edited  # pick first
+    if isinstance(item, str):
+        return item  # already a path
+    if isinstance(item, np.ndarray):
+        item = PIL.Image.fromarray(item)
+    assert isinstance(item, 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)
+    item.save(out_path)
     return out_path
 
 
@@ -288,4 +294,4 @@ if __name__ == "__main__":
             with gr.Tab(label="PartEdit", id="edit"):
                 edit_demo(model)
 
-        demo.queue(max_size=20).launch()
+        demo.queue(concurrency_count=1, max_size=20).launch(server_name="0.0.0.0")

requirements.txt

@@ -4,6 +4,7 @@
 
 setuptools>=61.0
 numpy<1.24
+# Not needed for demo environment
 # ipywidgets
 # black[jupyter]
 # jupyterlab
@@ -16,7 +17,7 @@ tqdm
 # Core ML stack (PyTorch 2.1.0 + CUDA 11.8)
 torch==2.1.0
 torchvision==0.16.0
-# torchaudio==2.1.0
+# torchaudio==2.1.0 
 
 # UI / HF stack
 gradio<5.0           # tested on 4.29; should work on 4.44.1 with pydantic fix

stable_diffusion_xl_partedit.py

@@ -1132,9 +1132,6 @@ class DotDictExtra(dict):
             if self.grounding.max() > 1.0:
                 self.grounding = self.grounding / self.grounding.max()
 
-            # self.edit_mask = ToTensor()(self.edit_mask.convert('')).unsqueeze(0)
-            # # TODO(Alex): Fix this
-            # self.edit_mask = torch.load(self.edit_mask)
         assert isinstance(self.th_strategy, Binarization), "th_strategy should be of type Binarization"
         assert isinstance(self.pad_strategy, PaddingStrategy), "pad_strategy should be of type PaddingStrategy"
 
@@ -1660,7 +1657,7 @@ class AttentionStore(AttentionControl):
 class AttentionControlEdit(AttentionStore, abc.ABC):
     def step_callback(self, x_t):
         if self.local_blend is not None:
-            # x_t = self.local_blend(x_t, self.attention_store) # TODO:return back
+            # x_t = self.local_blend(x_t, self.attention_store) # TODO: Check if there is more memory efficient way
             x_t = self.local_blend(x_t, self)
         return x_t
 
@@ -1688,7 +1685,7 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
             if is_cross:
                 alpha_words = self.cross_replace_alpha[self.cur_step].to(attn_base.device)
                 attn_replace_new = self.replace_cross_attention(attn_base, attn_replace) * alpha_words + (1 - alpha_words) * attn_replace
-                # TODO(Alex): We want to control the actual
+                
 
                 attn[1:] = attn_replace_new
                 if self.has_maps() and self.extra_kwargs.get("force_cross_attn", False):  # and self.cur_step <= 51: