Update app.py

app.py

@@ -187,7 +187,7 @@ class LLMService:
 
         model_id_or_path = "stablediffusionapi/realistic-vision-v51"
         self.vae_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, safety_checker=None, torch_dtype=torch.float16)
-        self.vae_pipe = self.vae_pipe.to("cpu")
+        self.vae_pipe = self.vae_pipe.to(self.vit_sd_device)
 
         self.boi_token_id = self.tokenizer.encode(BOI_TOKEN, add_special_tokens=False)[0]
         self.eoi_token_id = self.tokenizer.encode(EOI_TOKEN, add_special_tokens=False)[0]
@@ -361,8 +361,10 @@ def generate(text_list, image_list, max_new_tokens, force_boi, force_bbox, force
             generated_image = service.sd_adapter.generate(image_embeds=img_feat, num_inference_steps=50)[0]
 
             if force_polish:
-                service.sd_adapter = service.sd_adapter.cpu()
-                service.vae_pipe = service.vae_pipe.to(service.vit_sd_device, dtype=service.dtype)
+                #service.sd_adapter = service.sd_adapter.cpu()
+                #service.vae_pipe = service.vae_pipe.to(service.vit_sd_device, dtype=service.dtype)
+
+                torch.cuda.empty_cache()
                 
                 init_image = generated_image.resize((1024, 1024))
                 prompt = ""
@@ -373,8 +375,10 @@ def generate(text_list, image_list, max_new_tokens, force_boi, force_bbox, force
                 image_base64 = encode_image(generated_image)
                 gen_imgs_base64_list.append(image_base64)
 
-                service.vae_pipe = service.vae_pipe.to("cpu")
-                service.sd_adapter = service.sd_adapter.to(service.vit_sd_device, dtype=service.dtype)
+                torch.cuda.empty_cache()
+
+                # service.vae_pipe = service.vae_pipe.to("cpu")
+                # service.sd_adapter = service.sd_adapter.to(service.vit_sd_device, dtype=service.dtype)
                 
 
         # print('loading visual encoder and llm to GPU, and sd to CPU')