Update inferencer.py

inferencer.py

@@ -80,7 +80,7 @@ class Inferencer:
         self.image_size = image_size
         self.image_shape = (image_size // 16, image_size // 16)
         self.cfg_prompt = cfg_prompt
-        self.model = self.init_model()
+        self.model = None
 
     @spaces.GPU 
     def init_model(self):
@@ -132,6 +132,8 @@ class Inferencer:
         cfg_schedule="constant",
         temperature=1.0,
     ):
+        if not model:
+            self.model = self.init_model()
         prompt = self.model.prompt_template["INSTRUCTION"].format(
             input=f"Generate an image: {raw_prompt.strip()}."
         )
@@ -168,6 +170,8 @@ class Inferencer:
     @spaces.GPU 
     def query_image(self, img: Image.Image, prompt=""):
         model = self.model
+        if not model:
+            self.model = self.init_model()
         tokenizer = model.tokenizer
         special_tokens_dict = {"additional_special_tokens": ["<image>"]}
         tokenizer.add_special_tokens(special_tokens_dict)
@@ -214,143 +218,6 @@ class Inferencer:
         # print(tokenizer.decode(output[0]))
         return tokenizer.decode(output[0])
 
-    # def edit_image(
-    #     self,
-    #     img: Image.Image,
-    #     prompt: str,
-    #     cfg: float = 2.0,
-    #     cfg_prompt: str = "Repeat this image.",
-    #     cfg_schedule="constant",
-    #     temperature: float = 1.0,
-    #     grid_size: int = 2,
-    #     num_iter: int = 64,
-    #     mode: str = "conditional",
-    # ) -> list[Image.Image]:
-
-    #     model = self.model
-    #     tokenizer = model.tokenizer
-    #     m = n = self.image_size // 16
-    #     image_length = m * n + 64
-
-    #     # preprocess image
-    #     image = img.convert("RGB")
-    #     original_size = image.size
-    #     image = image.resize((self.image_size, self.image_size))
-    #     image = torch.from_numpy(np.array(image)).to(
-    #         dtype=model.dtype, device=self.device
-    #     )
-    #     image = rearrange(image, "h w c -> c h w")[None]
-    #     image = 2 * (image / 255) - 1
-
-    #     # prepare prompt
-    #     special_tokens_dict = {"additional_special_tokens": ["<image>"]}
-    #     tokenizer.add_special_tokens(special_tokens_dict)
-    #     image_token_idx = tokenizer.encode("<image>", add_special_tokens=False)[-1]
-
-    #     full_prompt = model.prompt_template["INSTRUCTION"].format(
-    #         input="<image>\n" + prompt
-    #     )
-    #     full_prompt = full_prompt.replace("<image>", "<image>" * image_length)
-    #     input_ids = tokenizer.encode(
-    #         full_prompt, add_special_tokens=True, return_tensors="pt"
-    #     )[0].to(self.device)
-
-    #     if cfg != 1.0:
-    #         null_prompt = model.prompt_template["INSTRUCTION"].format(
-    #             input="<image>\n" + cfg_prompt
-    #         )
-    #         null_prompt = null_prompt.replace("<image>", "<image>" * image_length)
-    #         null_input_ids = tokenizer.encode(
-    #             null_prompt, add_special_tokens=True, return_tensors="pt"
-    #         )[0].to(self.device)
-    #         attention_mask = pad_sequence(
-    #             [torch.ones_like(input_ids), torch.ones_like(null_input_ids)],
-    #             batch_first=True,
-    #             padding_value=0,
-    #         ).to(torch.bool)
-    #         input_ids = pad_sequence(
-    #             [input_ids, null_input_ids],
-    #             batch_first=True,
-    #             padding_value=tokenizer.eos_token_id,
-    #         )
-    #     else:
-    #         input_ids = input_ids[None]
-    #         attention_mask = torch.ones_like(input_ids).to(torch.bool)
-
-    #     with torch.no_grad():
-    #         x_enc = model.encode(image).to(model.dtype)
-    #         x_con, z_enc = model.extract_visual_feature(x_enc)
-
-    #     if cfg != 1.0:
-    #         z_enc = torch.cat([z_enc, z_enc], dim=0)
-    #         x_con = torch.cat([x_con, x_con], dim=0)
-
-    #     inputs_embeds = z_enc.new_zeros(*input_ids.shape, model.llm.config.hidden_size)
-    #     inputs_embeds[input_ids == image_token_idx] = z_enc.flatten(0, 1)
-    #     inputs_embeds[input_ids != image_token_idx] = model.llm.get_input_embeddings()(
-    #         input_ids[input_ids != image_token_idx]
-    #     )
-
-    #     # repeat
-    #     bsz = grid_size**2
-    #     x_con = torch.cat([x_con] * bsz)
-    #     if cfg != 1.0:
-    #         inputs_embeds = torch.cat(
-    #             [
-    #                 inputs_embeds[:1].expand(bsz, -1, -1),
-    #                 inputs_embeds[1:].expand(bsz, -1, -1),
-    #             ]
-    #         )
-    #         attention_mask = torch.cat(
-    #             [
-    #                 attention_mask[:1].expand(bsz, -1),
-    #                 attention_mask[1:].expand(bsz, -1),
-    #             ]
-    #         )
-    #     else:
-    #         inputs_embeds = inputs_embeds.expand(bsz, -1, -1)
-    #         attention_mask = attention_mask.expand(bsz, -1)
-
-    #     # sample
-    #     with torch.no_grad():
-    #         if mode == "conditional":
-    #             samples = model.sample(
-    #                 inputs_embeds=inputs_embeds,
-    #                 attention_mask=attention_mask,
-    #                 num_iter=num_iter,
-    #                 cfg=cfg,
-    #                 cfg_schedule=cfg_schedule,
-    #                 temperature=temperature,
-    #                 progress=False,
-    #                 image_shape=(m, n),
-    #                 x_con=x_con,
-    #             )
-    #         else:
-    #             samples = model.sample(
-    #                 inputs_embeds=inputs_embeds,
-    #                 attention_mask=attention_mask,
-    #                 num_iter=num_iter,
-    #                 cfg=cfg,
-    #                 cfg_schedule=cfg_schedule,
-    #                 temperature=temperature,
-    #                 progress=False,
-    #                 image_shape=(m, n),
-    #             )
-
-    #     samples = rearrange(
-    #         samples, "(m n) c h w -> (m h) (n w) c", m=grid_size, n=grid_size
-    #     )
-    #     samples = (
-    #         torch.clamp(127.5 * samples + 128.0, 0, 255)
-    #         .to("cpu", dtype=torch.uint8)
-    #         .numpy()
-    #     )
-
-    #     output_image = Image.fromarray(samples).resize(
-    #         (original_size[0] * grid_size, original_size[1] * grid_size)
-    #     )
-    #     return [output_image]
-
     @spaces.GPU 
     def edit_image(
         self, 
@@ -366,6 +233,8 @@ class Inferencer:
         """Edit single image based on prompt."""
 
         model = self.model
+        if not model:
+            self.model = self.init_model()
         tokenizer = model.tokenizer
         special_tokens_dict = {"additional_special_tokens": ["<image>"]}
         tokenizer.add_special_tokens(special_tokens_dict)
@@ -466,6 +335,8 @@ class Inferencer:
     @spaces.GPU 
     def query_text(self, prompt=""):
         model = self.model
+        if not model:
+            self.model = self.init_model()
         tokenizer = model.tokenizer
 
         # 构造文本 prompt