Update app.py

app.py

@@ -159,19 +159,27 @@ class LLMService:
         vae = AutoencoderKL.from_pretrained(args.diffusion_path, subfolder="vae").to(self.vit_sd_device, dtype=self.dtype)
 
 
-        unet = UNet2DConditionModel.from_pretrained(args.diffusion_path, subfolder="unet").to(dtype=self.dtype)
+        unet = UNet2DConditionModel.from_pretrained(args.diffusion_path, subfolder="unet").to(self.vit_sd_device, dtype=self.dtype)
 
         sd_adapter_cfg = OmegaConf.load(args.sd_adapter)
 
-        self.sd_adapter = hydra.utils.instantiate(sd_adapter_cfg, unet=unet).eval().to(dtype=self.dtype)
+        self.sd_adapter = hydra.utils.instantiate(sd_adapter_cfg, unet=unet).eval().to(self.vit_sd_device, dtype=self.dtype)
+
+        # self.sd_adapter.init_pipe(vae=vae,
+        #                           scheduler=noise_scheduler,
+        #                           visual_encoder=self.visual_encoder.cpu(),
+        #                           image_transform=self.image_transform,
+        #                           discrete_model=None,
+        #                           dtype=self.dtype,
+        #                           device="cpu")
 
         self.sd_adapter.init_pipe(vae=vae,
                                   scheduler=noise_scheduler,
-                                  visual_encoder=self.visual_encoder.cpu(),
+                                  visual_encoder=self.visual_encoder,
                                   image_transform=self.image_transform,
                                   discrete_model=None,
                                   dtype=self.dtype,
-                                  device="cpu")
+                                  device=self.vit_sd_device)
 
         print('Init sd adapter pipe done.')
 
@@ -336,11 +344,11 @@ def generate(text_list, image_list, max_new_tokens, force_boi, force_bbox):
     generated_text = generated_text.replace(EOI_TOKEN, IMG_FLAG).replace(service.tokenizer.eos_token, '')
 
     if output['has_img_output']:
-        print('loading visual encoder and llm to CPU, and sd to GPU')
-        a = time.time()
-        service.agent = service.agent.cpu()
-        service.sd_adapter = service.sd_adapter.to(service.vit_sd_device, dtype=service.dtype)
-        print("Loading finished: ", time.time() - a)
+        # print('loading visual encoder and llm to CPU, and sd to GPU')
+        # a = time.time()
+        # service.agent = service.agent.cpu()
+        # service.sd_adapter = service.sd_adapter.to(service.vit_sd_device, dtype=service.dtype)
+        # print("Loading finished: ", time.time() - a)
 
         img_gen_feat = output['img_gen_feat'].to(service.vit_sd_device, dtype=service.dtype)
 
@@ -350,12 +358,12 @@ def generate(text_list, image_list, max_new_tokens, force_boi, force_bbox):
             image_base64 = encode_image(generated_image)
             gen_imgs_base64_list.append(image_base64)
 
-        print('loading visual encoder and llm to GPU, and sd to CPU')
-        a = time.time()
-        service.sd_adapter = service.sd_adapter.cpu()
-        service.visual_encoder = service.visual_encoder.to(service.vit_sd_device, dtype=service.dtype)
-        service.agent = service.agent.to(service.vit_sd_device, dtype=service.dtype)
-        print("Loading finished: ", time.time() - a)
+        # print('loading visual encoder and llm to GPU, and sd to CPU')
+        # a = time.time()
+        # service.sd_adapter = service.sd_adapter.cpu()
+        # service.visual_encoder = service.visual_encoder.to(service.vit_sd_device, dtype=service.dtype)
+        # service.agent = service.agent.to(service.vit_sd_device, dtype=service.dtype)
+        # print("Loading finished: ", time.time() - a)
 
     if args.has_bbox:
         bboxes = extract_box(generated_text)