Download model locally

app.py

@@ -8,13 +8,12 @@ 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, unregister_attention_editor_diffusers
+from injection_utils import register_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"
+REMOTE_MODEL_PATH = "stabilityai/stable-diffusion-xl-base-1.0"
+LOCAL_MODEL_PATH = "./model"
 RESOLUTION = 256
 MIN_SIZE = 0.01
 WHITE = 255
@@ -114,7 +113,6 @@ FOOTNOTE = """
 
 
 def inference(
-    model,
     boxes,
     prompts,
     subject_token_indices,
@@ -135,7 +133,10 @@ def inference(
         raise gr.Error("cuda is not available")
 
     device = torch.device("cuda")
-    model.to(device).half()
+    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(LOCAL_MODEL_PATH, scheduler=scheduler, torch_dtype=torch.float16).to(device)
+    model.unet.set_attn_processor(AttnProcessor2_0())
+    model.enable_sequential_cpu_offload()
 
     seed_everything(seed)
     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
@@ -160,10 +161,7 @@ def inference(
     )
 
     register_attention_editor_diffusers(model, editor)
-    images = model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
-    unregister_attention_editor_diffusers(model)
-    model.double().to(torch.device("cpu"))
-    return images
+    return model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
 
 
 @spaces.GPU(duration=300)
@@ -198,7 +196,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)
 
@@ -255,10 +253,9 @@ 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)
-    model.unet.set_attn_processor(AttnProcessor2_0())
-    model.enable_sequential_cpu_offload()
+    model = StableDiffusionXLPipeline.from_pretrained(REMOTE_MODEL_PATH, scheduler=scheduler)
+    model.save_pretrained(LOCAL_MODEL_PATH)
+    del model
     
     with gr.Blocks(
             css=CSS,
@@ -330,7 +327,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,