Upload 18 files

modeling_vlm.py

@@ -27,12 +27,14 @@ from transformers import (
     PreTrainedModel,
     GenerationMixin
 )
+import numpy as np
 from transformers.configuration_utils import PretrainedConfig
 
 from .clip_encoder import CLIPVisionTower
 from .siglip_vit import create_siglip_vit
 from .projector import MlpProjector
 from .configuration_vlm import AttrDict, MultiModalityConfig, VisionConfig, AlignerConfig, GenVisionConfig, GenHeadConfig,  GenAlignerConfig
+from .vq_model import VQ_models
 
 
 class vision_head(torch.nn.Module):
@@ -61,7 +63,7 @@ def model_name_to_cls(cls_name):
         cls = CLIPVisionTower
 
     elif "VQ" in cls_name:
-        from janus.models.vq_model import VQ_models
+        from .vq_model import VQ_models
 
         cls = VQ_models[cls_name]
     elif "vision_head" in cls_name:
@@ -193,7 +195,87 @@ class MultiModalityCausalLM(MultiModalityPreTrainedModel):
             inputs_embeds = self.prepare_inputs_embeds(input_ids, pixel_values, images_seq_mask, images_emb_mask, **kwargs)
         return self.language_model.generate(inputs_embeds=inputs_embeds, past_key_values=past_key_values, attention_mask=attention_mask, position_ids=position_ids, **kwargs)
 
+    @torch.no_grad()    
+    def generate_image(
+        self,
+        processor,
+        prompt: str,
+        temperature: float = 1,
+        parallel_size: int = 16,
+        cfg_weight: float = 5,
+        image_token_num_per_image: int = 576,
+        img_size: int = 384,
+        patch_size: int = 16,
+        generator=None
+    ):
+        from PIL import Image
+
+        conversation = [
+            {
+                "role": "User",
+                "content": prompt,
+            },
+            {"role": "Assistant", "content": ""},
+        ]
+
+        sft_format = processor.apply_sft_template_for_multi_turn_prompts(
+            conversations=conversation,
+            sft_format=processor.sft_format,
+            system_prompt="",
+        )
+        prompt = sft_format + processor.image_start_tag
+        input_ids = processor.tokenizer.encode(prompt)
+        input_ids = torch.LongTensor(input_ids)
+
+        tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int)
+        for i in range(parallel_size * 2):
+            tokens[i, :] = input_ids
+            if i % 2 != 0:
+                tokens[i, 1:-1] = processor.pad_id
+
+        inputs_embeds = self.language_model.get_input_embeddings()(tokens)
+
+        generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int)
+        past_key_values = None
+
+        for i in range(image_token_num_per_image):
+            outputs = self.language_model.model.forward(
+                input_ids=None,
+                inputs_embeds=inputs_embeds,
+                use_cache=True,
+                past_key_values=past_key_values,
+            )
+            hidden_states = outputs.last_hidden_state
+            past_key_values = outputs.past_key_values
+            logits = self.gen_head(hidden_states[:, -1, :])
+            logit_cond = logits[0::2, :]
+            logit_uncond = logits[1::2, :]
+
+            logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
+            probs = torch.softmax(logits / temperature, dim=-1)
+
+            next_token = torch.multinomial(probs, num_samples=1) if generator is None else torch.multinomial(probs, num_samples=1, generator=generator)
+            generated_tokens[:, i] = next_token.squeeze(dim=-1)
+
+            next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
+            img_embeds = self.prepare_gen_img_embeds(next_token)
+            inputs_embeds = img_embeds.unsqueeze(dim=1)
+        dec = self.gen_vision_model.decode_code(
+            generated_tokens.to(dtype=torch.int), [parallel_size, 8, img_size // patch_size, img_size // patch_size]
+        )
+        dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
+
+        dec = np.clip((dec + 1) / 2 * 255, 0, 255)
+
+        visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
+        visual_img[:, :, :] = dec
+
+        images = []
+
+        for i in range(parallel_size):
+            images.append(Image.fromarray(visual_img[i]))
 
+        return images
 
 
 AutoConfig.register("vision", VisionConfig)