Use diffusers implementation

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "Attend-and-Excite"]
-	path = Attend-and-Excite
-	url = https://github.com/AttendAndExcite/Attend-and-Excite

.vscode/settings.json

@@ -0,0 +1,18 @@
+{
+    "python.linting.enabled": true,
+    "python.linting.flake8Enabled": true,
+    "python.linting.pylintEnabled": false,
+    "python.linting.lintOnSave": true,
+    "python.formatting.provider": "yapf",
+    "python.formatting.yapfArgs": [
+        "--style={based_on_style: pep8, indent_width: 4, blank_line_before_nested_class_or_def: false, spaces_before_comment: 2, split_before_logical_operator: true}"
+    ],
+    "[python]": {
+        "editor.formatOnType": true,
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": true
+        }
+    },
+    "editor.formatOnSave": true,
+    "files.insertFinalNewline": true
+}

Attend-and-Excite

@@ -1 +0,0 @@
-Subproject commit 41620338367f980b9d73752360ffd2557d8ddf5d

app.py

@@ -24,12 +24,11 @@ def process_example(
     seed: int,
     apply_attend_and_excite: bool,
 ) -> tuple[list[tuple[int, str]], PIL.Image.Image]:
-    model_id = 'CompVis/stable-diffusion-v1-4'
     num_steps = 50
     guidance_scale = 7.5
 
-    token_table = model.get_token_table(model_id, prompt)
-    result = model.run(model_id, prompt, indices_to_alter_str, seed,
+    token_table = model.get_token_table(prompt)
+    result = model.run(prompt, indices_to_alter_str, seed,
                        apply_attend_and_excite, num_steps, guidance_scale)
     return token_table, result
 
@@ -39,9 +38,6 @@ with gr.Blocks(css='style.css') as demo:
 
     with gr.Row():
         with gr.Column():
-            model_id = gr.Text(label='Model ID',
-                               value='CompVis/stable-diffusion-v1-4',
-                               visible=False)
             prompt = gr.Text(
                 label='Prompt',
                 max_lines=1,
@@ -171,13 +167,12 @@ with gr.Blocks(css='style.css') as demo:
 
     show_token_indices_button.click(
         fn=model.get_token_table,
-        inputs=[model_id, prompt],
+        inputs=prompt,
         outputs=token_indices_table,
         queue=False,
     )
 
     inputs = [
-        model_id,
         prompt,
         token_indices_str,
         seed,
@@ -187,7 +182,7 @@ with gr.Blocks(css='style.css') as demo:
     ]
     prompt.submit(
         fn=model.get_token_table,
-        inputs=[model_id, prompt],
+        inputs=prompt,
         outputs=token_indices_table,
         queue=False,
     ).then(
@@ -197,7 +192,7 @@ with gr.Blocks(css='style.css') as demo:
     )
     token_indices_str.submit(
         fn=model.get_token_table,
-        inputs=[model_id, prompt],
+        inputs=prompt,
         outputs=token_indices_table,
         queue=False,
     ).then(
@@ -207,7 +202,7 @@ with gr.Blocks(css='style.css') as demo:
     )
     run_button.click(
         fn=model.get_token_table,
-        inputs=[model_id, prompt],
+        inputs=prompt,
         outputs=token_indices_table,
         queue=False,
     ).then(

model.py

@@ -1,83 +1,73 @@
 from __future__ import annotations
 
-import sys
-
-import gradio as gr
 import PIL.Image
 import torch
-
-sys.path.append('Attend-and-Excite')
-
-from config import RunConfig
-from pipeline_attend_and_excite import AttendAndExcitePipeline
-from run import run_on_prompt
-from utils.ptp_utils import AttentionStore
+from diffusers import (StableDiffusionAttendAndExcitePipeline,
+                       StableDiffusionPipeline)
 
 
 class Model:
     def __init__(self):
         self.device = torch.device(
             'cuda:0' if torch.cuda.is_available() else 'cpu')
-        self.model_id = ''
-        self.model = None
-        self.tokenizer = None
-
-        self.load_model('CompVis/stable-diffusion-v1-4')
+        model_id = 'CompVis/stable-diffusion-v1-4'
+        if self.device.type == 'cuda':
+            self.ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(
+                model_id, torch_dtype=torch.float16)
+            self.ax_pipe.to(self.device)
+            self.sd_pipe = StableDiffusionPipeline.from_pretrained(
+                model_id, torch_dtype=torch.float16)
+            self.sd_pipe.to(self.device)
+        else:
+            self.ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(
+                model_id)
+            self.sd_pipe = StableDiffusionPipeline.from_pretrained(model_id)
 
-    def load_model(self, model_id: str) -> None:
-        if model_id == self.model_id:
-            return
-        self.model = AttendAndExcitePipeline.from_pretrained(model_id).to(
-            self.device)
-        self.tokenizer = self.model.tokenizer
-        self.model_id = model_id
-
-    def get_token_table(self, model_id: str, prompt: str):
-        self.load_model(model_id)
+    def get_token_table(self, prompt: str):
         tokens = [
-            self.tokenizer.decode(t)
-            for t in self.tokenizer(prompt)['input_ids']
+            self.ax_pipe.tokenizer.decode(t)
+            for t in self.ax_pipe.tokenizer(prompt)['input_ids']
         ]
         tokens = tokens[1:-1]
         return list(enumerate(tokens, start=1))
 
     def run(
         self,
-        model_id: str,
         prompt: str,
         indices_to_alter_str: str,
-        seed: int,
-        apply_attend_and_excite: bool,
-        num_steps: int,
-        guidance_scale: float,
+        seed: int = 0,
+        apply_attend_and_excite: bool = True,
+        num_steps: int = 50,
+        guidance_scale: float = 7.5,
         scale_factor: int = 20,
         thresholds: dict[int, float] = {
             10: 0.5,
-            20: 0.8
+            20: 0.8,
         },
         max_iter_to_alter: int = 25,
     ) -> PIL.Image.Image:
         generator = torch.Generator(device=self.device).manual_seed(seed)
-        try:
-            indices_to_alter = list(map(int, indices_to_alter_str.split(',')))
-        except:
-            raise gr.Error('Invalid token indices.')
-
-        self.load_model(model_id)
-
-        controller = AttentionStore()
-        config = RunConfig(prompt=prompt,
-                           n_inference_steps=num_steps,
-                           guidance_scale=guidance_scale,
-                           run_standard_sd=not apply_attend_and_excite,
-                           scale_factor=scale_factor,
-                           thresholds=thresholds,
-                           max_iter_to_alter=max_iter_to_alter)
-        image = run_on_prompt(model=self.model,
-                              prompt=[prompt],
-                              controller=controller,
-                              token_indices=indices_to_alter,
-                              seed=generator,
-                              config=config)
 
-        return image
+        if apply_attend_and_excite:
+            try:
+                token_indices = list(map(int, indices_to_alter_str.split(',')))
+            except Exception:
+                raise ValueError('Invalid token indices.')
+            out = self.ax_pipe(
+                prompt=prompt,
+                token_indices=token_indices,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                num_inference_steps=num_steps,
+                max_iter_to_alter=max_iter_to_alter,
+                thresholds=thresholds,
+                scale_factor=scale_factor,
+            )
+        else:
+            out = self.sd_pipe(
+                prompt=prompt,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                num_inference_steps=num_steps,
+            )
+        return out.images[0]

requirements.txt

@@ -1,8 +1,5 @@
-accelerate==0.19.0
-diffusers==0.12.1
-ftfy==6.1.1
-jupyter
-opencv-python-headless==4.7.0.68
-pyrallis==0.3.1
-torch==1.13.1
-transformers==4.29.2
+accelerate==0.20.3
+diffusers==0.17.0
+Pillow==9.5.0
+torch==2.0.1
+transformers==4.30.1