Creates model inside GPU

app.py

@@ -10,7 +10,6 @@ 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
 
@@ -21,7 +20,6 @@ COLORS = ["red", "blue", "green", "orange", "purple", "turquoise", "olive"]
 
 
 def inference(
-    model,
     boxes,
     prompts,
     subject_token_indices,
@@ -42,55 +40,41 @@ def inference(
         raise gr.Error("cuda is not available")
 
     device = torch.device("cuda")
-    model = model.to(device=device, dtype=torch.float16)
+    model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+    model = StableDiffusionXLPipeline.from_pretrained(model_path, scheduler=scheduler, device=device, torch_dtype=torch.float16)
+    model.unet.set_default_attn_processor()
+    model.enable_sequential_cpu_offload()
 
     seed_everything(seed)
     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
     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)
+    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
 
 
 @spaces.GPU
 def generate(
-    model,
     prompt,
     subject_token_indices,
     filter_token_indices,
@@ -120,7 +104,7 @@ def generate(
     prompts = [prompt.strip('.').strip(',').strip()] * batch_size
 
     images = inference(
-        model, boxes, prompts, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
+        boxes, prompts, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
         final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale, classifier_free_guidance_scale,
         num_iterations, loss_threshold, num_guidance_steps, seed)
 
@@ -214,10 +198,6 @@ def main():
     }
     """
 
-    model_path = "stabilityai/stable-diffusion-xl-base-1.0"
-    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
-    model = StableDiffusionXLPipeline.from_pretrained(model_path, scheduler=scheduler)
-
     nltk.download('averaged_perceptron_tagger')
 
     with gr.Blocks(
@@ -328,7 +308,7 @@ def main():
             )
 
             generate_image_button.click(
-                fn=partial(generate, model),
+                fn=generate,
                 inputs=[
                     prompt, subject_token_indices, filter_token_indices, num_tokens,
                     init_step_size, final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale,