Load model to CPU

app.py

@@ -6,11 +6,14 @@ import numpy as np
 from PIL import Image, ImageDraw
 
 from diffusers import DDIMScheduler
+from diffusers.models.attention_processor import AttnProcessor2_0
 from pipeline_stable_diffusion_xl_opt import StableDiffusionXLPipeline
-from injection_utils import register_attention_editor_diffusers
+from injection_utils import register_attention_editor_diffusers, unregister_attention_editor_diffusers
 from bounded_attention import BoundedAttention
 from pytorch_lightning import seed_everything
 
+from functools import partial
+
 MODEL_PATH = "stabilityai/stable-diffusion-xl-base-1.0"
 RESOLUTION = 256
 MIN_SIZE = 0.01
@@ -111,6 +114,7 @@ FOOTNOTE = """
 
 
 def inference(
+    model,
     boxes,
     prompts,
     subject_token_indices,
@@ -131,10 +135,7 @@ def inference(
         raise gr.Error("cuda is not available")
 
     device = torch.device("cuda")
-    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, torch_dtype=torch.float16).to(device)
-    model.unet.set_default_attn_processor()
-    model.enable_sequential_cpu_offload()
+    model.to(device)
 
     seed_everything(seed)
     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
@@ -159,11 +160,15 @@ def inference(
     )
 
     register_attention_editor_diffusers(model, editor)
-    return model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
+    images = model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
+    unregister_attention_editor_diffusers(model)
+    model.to(torch.device("cpu"))
+    return images
 
 
 @spaces.GPU(duration=300)
 def generate(
+    model,
     prompt,
     subject_token_indices,
     filter_token_indices,
@@ -193,7 +198,7 @@ def generate(
     prompts = [prompt.strip(".").strip(",").strip()] * batch_size
 
     images = inference(
-        boxes, prompts, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
+        model, 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)
 
@@ -249,6 +254,11 @@ def clear(batch_size):
 
 def main():
     nltk.download("averaged_perceptron_tagger")
+
+    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, torch_dtype=torch.float16).to(device)
+    model.unet.set_attn_processor(AttnProcessor2_0())
+    model.enable_sequential_cpu_offload()
     
     with gr.Blocks(
             css=CSS,
@@ -320,7 +330,7 @@ def main():
             )
     
             generate_image_button.click(
-                fn=generate,
+                fn=partial(generate, model),
                 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,