Update app.py

app.py

@@ -84,6 +84,7 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
             attention_mask=prompt_masks,
             output_hidden_states=True,
         ).hidden_states[-2]
+        text_encoder.cpu()
 
     return prompt_embeds, prompt_masks
 
@@ -242,13 +243,13 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
                     torch.random.manual_seed(int(seed))
                 z = torch.randn([1, 16, latent_h, latent_w], device="cuda").to(dtype)
                 z = z.repeat(2, 1, 1, 1)
-
+                model.cpu()
                 with torch.no_grad():
                     if neg_cap != "":
                         cap_feats, cap_mask = encode_prompt([cap] + [neg_cap], text_encoder, tokenizer, 0.0)
                     else:
                         cap_feats, cap_mask = encode_prompt([cap] + [""], text_encoder, tokenizer, 0.0)
-
+                model.cuda()
                 cap_mask = cap_mask.to(cap_feats.device)
 
                 model_kwargs = dict(