Update

app.py

@@ -6,8 +6,8 @@ import os
 
 import gradio as gr
 import PIL.Image
-
-from model import Model
+import torch
+from diffusers import StableDiffusionAttendAndExcitePipeline, StableDiffusionPipeline
 
 DESCRIPTION = """\
 # Attend-and-Excite
@@ -17,7 +17,63 @@ Attend-and-Excite performs attention-based generative semantic guidance to mitig
 Select a prompt and a set of indices matching the subjects you wish to strengthen (the `Check token indices` cell can help map between a word and its index).
 """
 
-model = Model()
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+if torch.cuda.is_available():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model_id = "CompVis/stable-diffusion-v1-4"
+    ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(model_id)
+    ax_pipe.to(device)
+    sd_pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    sd_pipe.to(device)
+
+
+def get_token_table(prompt: str) -> list[tuple[int, str]]:
+    tokens = [ax_pipe.tokenizer.decode(t) for t in ax_pipe.tokenizer(prompt)["input_ids"]]
+    tokens = tokens[1:-1]
+    return list(enumerate(tokens, start=1))
+
+
+def run(
+    prompt: str,
+    indices_to_alter_str: str,
+    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,
+    },
+    max_iter_to_alter: int = 25,
+) -> PIL.Image.Image:
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    if apply_attend_and_excite:
+        try:
+            token_indices = list(map(int, indices_to_alter_str.split(",")))
+        except Exception:
+            raise ValueError("Invalid token indices.")
+        out = 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 = sd_pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            num_inference_steps=num_steps,
+        )
+    return out.images[0]
 
 
 def process_example(
@@ -26,11 +82,13 @@ def process_example(
     seed: int,
     apply_attend_and_excite: bool,
 ) -> tuple[list[tuple[int, str]], PIL.Image.Image]:
-    num_steps = 50
-    guidance_scale = 7.5
-
-    token_table = model.get_token_table(prompt)
-    result = model.run(prompt, indices_to_alter_str, seed, apply_attend_and_excite, num_steps, guidance_scale)
+    token_table = get_token_table(prompt)
+    result = run(
+        prompt=prompt,
+        indices_to_alter_str=indices_to_alter_str,
+        seed=seed,
+        apply_attend_and_excite=apply_attend_and_excite,
+    )
     return token_table, result
 
 
@@ -176,11 +234,11 @@ with gr.Blocks(css="style.css") as demo:
         )
 
     show_token_indices_button.click(
-        fn=model.get_token_table,
+        fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
-        api_name=False,
+        api_name="get-token-table",
     )
 
     inputs = [
@@ -192,37 +250,37 @@ with gr.Blocks(css="style.css") as demo:
         guidance_scale,
     ]
     prompt.submit(
-        fn=model.get_token_table,
+        fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
         api_name=False,
     ).then(
-        fn=model.run,
+        fn=run,
         inputs=inputs,
         outputs=result,
         api_name=False,
     )
     token_indices_str.submit(
-        fn=model.get_token_table,
+        fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
         api_name=False,
     ).then(
-        fn=model.run,
+        fn=run,
         inputs=inputs,
         outputs=result,
         api_name=False,
     )
     run_button.click(
-        fn=model.get_token_table,
+        fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
         queue=False,
         api_name=False,
     ).then(
-        fn=model.run,
+        fn=run,
         inputs=inputs,
         outputs=result,
         api_name="run",

model.py

@@ -1,61 +0,0 @@
-from __future__ import annotations
-
-import PIL.Image
-import torch
-from diffusers import StableDiffusionAttendAndExcitePipeline, StableDiffusionPipeline
-
-
-class Model:
-    def __init__(self):
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        model_id = "CompVis/stable-diffusion-v1-4"
-        self.ax_pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained(model_id)
-        self.ax_pipe.to(self.device)
-        self.sd_pipe = StableDiffusionPipeline.from_pretrained(model_id)
-        self.sd_pipe.to(self.device)
-
-    def get_token_table(self, prompt: str):
-        tokens = [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,
-        prompt: str,
-        indices_to_alter_str: str,
-        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,
-        },
-        max_iter_to_alter: int = 25,
-    ) -> PIL.Image.Image:
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-
-        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]