Move model to GPU only inside generate function

app.py

@@ -21,7 +21,6 @@ COLORS = ["red", "blue", "green", "orange", "purple", "turquoise", "olive"]
 
 
 def inference(
-    device,
     model,
     boxes,
     prompts,
@@ -39,6 +38,12 @@ def inference(
     num_guidance_steps,
     seed,
 ):
+    if not torch.cuda.is_available():
+        raise gr.Error("cuda is not available")
+
+    device = torch.device("cuda")
+    model = model.to(device=device, dtype=torch.float16)
+
     seed_everything(seed)
     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
     eos_token_index = num_tokens + 1
@@ -85,7 +90,6 @@ def inference(
 
 @spaces.GPU
 def generate(
-    device,
     model,
     prompt,
     subject_token_indices,
@@ -107,8 +111,8 @@ def generate(
     subject_token_indices = convert_token_indices(subject_token_indices, nested=True)
     if len(boxes) != len(subject_token_indices):
         raise gr.Error("""
-            The number of boxes should be equal to the number of subject token indices.
-            Number of boxes drawn: {}, number of grounding tokens: {}.
+            The number of boxes should be equal to the number of subjects.
+            Number of boxes drawn: {}, number of subjects: {}.
         """.format(len(boxes), len(subject_token_indices)))
 
     filter_token_indices = convert_token_indices(filter_token_indices) if len(filter_token_indices.strip()) > 0 else None
@@ -116,7 +120,7 @@ def generate(
     prompts = [prompt.strip('.').strip(',').strip()] * batch_size
 
     images = inference(
-        device, model, 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)
 
@@ -210,10 +214,9 @@ def main():
     }
     """
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     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)
+    model = StableDiffusionXLPipeline.from_pretrained(model_path, scheduler=scheduler)
 
     nltk.download('averaged_perceptron_tagger')
 
@@ -325,7 +328,7 @@ def main():
             )
 
             generate_image_button.click(
-                fn=partial(generate, device, model),
+                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,

pipeline_stable_diffusion_xl_opt.py

@@ -831,6 +831,8 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
             num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
             timesteps = timesteps[:num_inference_steps]
 
+        latents = latents.half()
+        prompt_embeds = prompt_embeds.half()
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 latents = self.update_loss(latents, i, t, prompt_embeds, cross_attention_kwargs, add_text_embeds, add_time_ids)