feat: add web UI Demo with attention visualization and random mask generation

.gitignore

@@ -0,0 +1,105 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# IDE / Editor specific files
+.idea/
+.vscode/
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# pytype static analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

docs/README.md

@@ -215,6 +215,38 @@ Refer to the provided Jupyter notebooks for complete examples and results:
 ![Example Output](imgs/example.png)
 
 
+## 💻 Web UI Demo
+
+To make anomaly generation more interactive, we provide a Gradio-based web interface.
+
+**Installation:**
+```bash
+pip install gradio
+pip install huggingface_hub==0.24.5
+```
+
+**How to Launch:**
+```bash
+python web_demo.py
+```
+
+This will start a local web server. Open the provided URL in your browser to access the UI.
+
+**Features:**
+- **Interactive Prompting:** Dynamically create and preview prompts using templates.
+- **Real-time Generation:** Upload a reference image and generate anomaly variations on the fly.
+- **Flexible Masking:** Use your own mask or let the UI generate a random one for you.
+- **Advanced Controls:** Easily tweak parameters like guidance strength, attention scaling, and random seeds.
+- **Instant Visualization:** View the final image, attention maps, and generation metadata all in one place.
+
+*We recommend preparing an image like `imgs/web_ui_demo.png` to showcase your UI.*
+
+<details>
+<summary>Web UI Demo screenshot</summary>
+
+![Web UI Demo](imgs/web_ui_demo.png)
+</details>
+
 ## 🛠️ Todo List
 - [ ] Colab demo.
 - [ ] HuggingFace demo.

env.yml

@@ -141,7 +141,7 @@ dependencies:
       - bleach==6.1.0
       - blobfile==2.1.1
       - click==8.1.7
-      - clip==1.0
+      - clip==0.2.0
       - colorama==0.4.6
       - comm==0.2.2
       - contourpy==1.2.1
@@ -201,7 +201,6 @@ dependencies:
       - matplotlib-inline==0.1.7
       - mdurl==0.1.2
       - mistune==3.0.2
-      - mmcv==2.2.0
       - mmengine==0.10.4
       - model-index==0.1.11
       - multidict==6.0.5
@@ -217,10 +216,10 @@ dependencies:
       - notebook-shim==0.2.4
       - ogb==1.3.6
       - open-clip-torch==2.10.1
+      - open-clip==1.0.1
       - opencv-python==4.8.1.78
       - opendatalab==0.0.10
       - openmim==0.3.9
-      - openxlab==0.1.1
       - ordered-set==4.1.0
       - oss2==2.17.0
       - outdated==0.2.2
@@ -252,7 +251,6 @@ dependencies:
       - pyzmq==26.0.3
       - referencing==0.35.1
       - regex==2023.6.3
-      - requests==2.28.2
       - rfc3339-validator==0.1.4
       - rfc3986-validator==0.1.1
       - rich==13.4.2
@@ -278,7 +276,6 @@ dependencies:
       - tokenizers==0.13.3
       - tomli==2.0.1
       - tornado==6.4.1
-      - tqdm==4.65.2
       - traitlets==5.14.3
       - transformers==4.29.2
       - types-python-dateutil==2.9.0.20240316

run.py

@@ -1,6 +1,5 @@
 import pprint
-from typing import List
-
+from typing import List, Optional, Dict, Tuple, Callable
 import pyrallis
 import torch
 from PIL import Image
@@ -39,18 +38,33 @@ def get_indices_to_alter(stable, prompt: str) -> List[int]:
     return token_indices
 
 
-def get_indices_to_alter_new(stable, prompt: str, tokens) -> List[int]:
-    token_idx_to_word = {idx: stable.tokenizer.decode(t)
-                         for idx, t in enumerate(stable.tokenizer(prompt)['input_ids'])
-                         if 0 < idx < len(stable.tokenizer(prompt)['input_ids']) - 1}
+def get_indices_to_alter_new(stable, prompt: str, tokens_str: str) -> Tuple[List[int], Dict[int, str]]:
+    """
+    Parses the prompt to get a map of token indices to words,
+    and identifies the indices to alter based on user input.
+    Returns the list of indices to alter and the full token map.
+    """
+    token_ids = stable.tokenizer(prompt)['input_ids']
+    token_idx_to_word = {
+        idx: stable.tokenizer.decode(t)
+        for idx, t in enumerate(token_ids)
+        if 0 < idx < len(token_ids) - 1
+    }
     pprint.pprint(token_idx_to_word)
-    # token_indices = input("Please enter the a comma-separated list indices of the tokens you wish to "
-    #                       "alter (e.g., 2,5): ")
-    # pprint.pprint(token_indices)
-    token_indices = tokens
-    token_indices = [int(i) for i in token_indices.split(",")]
-    print(f"Altering tokens: {[token_idx_to_word[i] for i in token_indices]}")
-    return token_indices
+
+    # The web UI passes a number, which we convert to a string.
+    if isinstance(tokens_str, (int, float)):
+        tokens_str = str(int(tokens_str))
+
+    try:
+        # Handle comma-separated strings for indices
+        token_indices = [int(i.strip()) for i in tokens_str.split(",") if i.strip()]
+    except (ValueError, AttributeError) as e:
+        print(f"Warning: Could not parse token indices '{tokens_str}'. Error: {e}. Defaulting to empty list.")
+        token_indices = []
+
+    print(f"Altering tokens: {[token_idx_to_word.get(i, '<UNK>') for i in token_indices]}")
+    return token_indices, token_idx_to_word
 
 
 def run_on_prompt(prompt: List[str],
@@ -150,13 +164,13 @@ def run_on_prompt_and_masked_image(prompt: List[str],
                                    mask_image: str,
                                    seed: torch.Generator,
                                    config: RunConfig,
-
                                    normal_prompt,
                                    detailed_prompt,
-
-                                   img_prompt = None,
-                                   abnormal_img = None,
-                                   clip_loss = None) -> Image.Image:
+                                   img_prompt=None,
+                                   abnormal_img=None,
+                                   clip_loss=None,
+                                   callback: Optional[Callable] = None,
+                                   callback_steps: Optional[int] = 1) -> Image.Image:
     if controller is not None:
         ptp_utils.register_attention_control(model, controller)
     outputs, image_latents = model(prompt=prompt,
@@ -182,7 +196,9 @@ def run_on_prompt_and_masked_image(prompt: List[str],
                                    normal_prompt=normal_prompt,
                                    abnormal_img=abnormal_img,
                                    detailed_prompt=detailed_prompt,
-                                   clip_loss=clip_loss)
+                                   clip_loss=clip_loss,
+                                   callback=callback,
+                                   callback_steps=callback_steps)
     # image = outputs.images[0]
     # return image, image_latents
     return outputs[0], image_latents

utils/vis_utils.py

@@ -23,15 +23,20 @@ def show_cross_attention(prompt: str,
     attention_maps = aggregate_attention(attention_store, res, from_where, True, select).detach().cpu()
     images = []
 
-    # show spatial attention for indices of tokens to strengthen
+    # show spatial attention for all tokens
     for i in range(len(tokens)):
+        if i == 0 or i == len(tokens) - 1: # Skip start and end tokens
+            continue
         image = attention_maps[:, :, i]
+        image = show_image_relevance(image, orig_image)
+        image = image.astype(np.uint8)
+        image = np.array(Image.fromarray(image).resize((res ** 2, res ** 2)))
+        token_text = decoder(int(tokens[i]))
+        # Add a visual indicator for altered tokens
         if i in indices_to_alter:
-            image = show_image_relevance(image, orig_image)
-            image = image.astype(np.uint8)
-            image = np.array(Image.fromarray(image).resize((res ** 2, res ** 2)))
-            image = ptp_utils.text_under_image(image, decoder(int(tokens[i])))
-            images.append(image)
+            token_text = f"*{token_text}*"
+        image = ptp_utils.text_under_image(image, token_text)
+        images.append(image)
 
     ptp_utils.view_images(np.stack(images, axis=0))
     # TODO