Allow repeated inference

app.py

@@ -10,6 +10,7 @@ from pipeline_stable_diffusion_xl_opt import StableDiffusionXLPipeline
 from injection_utils import regiter_attention_editor_diffusers
 from bounded_attention import BoundedAttention
 from pytorch_lightning import seed_everything
+from torch_kmeans import KMeans
 
 from functools import partial
 
@@ -40,26 +41,45 @@ def inference(
 ):
     seed_everything(seed)
     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
-    editor = BoundedAttention(
-        boxes,
-        prompts,
-        subject_token_indices,
-        list(range(70, 82)),
-        list(range(70, 82)),
-        eos_token_index=num_tokens + 1,
-        cross_loss_coef=cross_loss_scale,
-        self_loss_coef=self_loss_scale,
-        filter_token_indices=filter_token_indices,
-        max_guidance_iter=num_guidance_steps,
-        max_guidance_iter_per_step=num_iterations,
-        start_step_size=init_step_size,
-        end_step_size=final_step_size,
-        loss_stopping_value=loss_threshold,
-        num_clusters_per_box=num_clusters_per_subject,
-        debug=False,
-    )
-
-    regiter_attention_editor_diffusers(model, editor)
+    eos_token_index = num_tokens + 1
+
+    if hasattr(model, 'editor'):
+        editor.boxes = boxes
+        editor.prompts = prompts
+        editor.subject_token_indices = subject_token_indices
+        editor.filter_token_indices = filter_token_indices
+        editor.eos_token_index = eos_token_index
+        editor.cross_loss_coef = cross_loss_scale
+        editor.self_loss_coef = self_loss_scale
+        editor.max_guidance_iter = num_guidance_steps
+        editor.max_guidance_iter_per_step = num_iterations
+        editor.start_step_size = init_step_size
+        self.step_size_coef = (final_step_size - init_step_size) / num_guidance_steps
+        editor.loss_stopping_value = loss_threshold
+        num_clusters = len(boxes) * num_clusters_per_subject
+        self.clustering = KMeans(n_clusters=num_clusters, num_init=100)
+
+    else:
+        editor = BoundedAttention(
+            boxes,
+            prompts,
+            subject_token_indices,
+            list(range(70, 82)),
+            list(range(70, 82)),
+            filter_token_indices=filter_token_indices,
+            eos_token_index=eos_token_index,
+            cross_loss_coef=cross_loss_scale,
+            self_loss_coef=self_loss_scale,
+            max_guidance_iter=num_guidance_steps,
+            max_guidance_iter_per_step=num_iterations,
+            start_step_size=init_step_size,
+            end_step_size=final_step_size,
+            loss_stopping_value=loss_threshold,
+            num_clusters_per_box=num_clusters_per_subject,
+        )
+
+        regiter_attention_editor_diffusers(model, editor)
+
     return model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images