from typing import List, Optional, Tuple, Union, Dict import torch import torch.nn as nn from PIL import Image import torch.nn.functional as F import transformers from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.generation.utils import GenerateOutput from blip3o.model.blip3o_arch import blip3oMetaModel, blip3oMetaForCausalLM from transformers import Qwen2_5_VLConfig, Qwen2_5_VLModel, Qwen2_5_VLForConditionalGeneration from blip3o.constants import UND_IMAGE_TOKEN_IDX from diffusers.utils.torch_utils import randn_tensor from diffusers.pipelines.pipeline_utils import numpy_to_pil import numpy as np from diffusers.models import AutoencoderKL from diffusers.schedulers import FlowMatchEulerDiscreteScheduler class blip3oQwenConfig(Qwen2_5_VLConfig): model_type = "blip3o_qwen" class blip3oQwenModel(blip3oMetaModel, Qwen2_5_VLModel): config_class = blip3oQwenConfig def __init__(self, config: Qwen2_5_VLConfig): super(blip3oQwenModel, self).__init__(config) class blip3oQwenForCausalLM(Qwen2_5_VLForConditionalGeneration, blip3oMetaForCausalLM): config_class = blip3oQwenConfig def __init__(self, config): Qwen2_5_VLForConditionalGeneration.__init__(self, config) config.model_type = "blip3o_qwen" self.model = blip3oQwenModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def get_model(self): return self.model # def forward( # self, # input_ids: torch.LongTensor = None, # attention_mask: Optional[torch.Tensor] = None, # position_ids: Optional[torch.LongTensor] = None, # past_key_values: Optional[List[torch.FloatTensor]] = None, # inputs_embeds: Optional[torch.FloatTensor] = None, # labels: Optional[torch.LongTensor] = None, # ids: Optional[list] = None, # i_s_pos: Optional[list] = None, # use_cache: Optional[bool] = None, # output_attentions: Optional[bool] = None, # output_hidden_states: Optional[bool] = None, # gen_image: Optional[torch.FloatTensor] = None, # und_image: Optional[torch.FloatTensor] = None, # grid_thw: Optional[torch.FloatTensor] = None, # image_sizes: Optional[List[List[int]]] = None, # return_dict: Optional[bool] = None, # cache_position: Optional[torch.LongTensor] = None # ) -> Union[Tuple, CausalLMOutputWithPast]: # output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions # output_hidden_states = ( # output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states # ) # return_dict = return_dict if return_dict is not None else self.config.use_return_dict # if inputs_embeds is None: # ( # input_ids, # position_ids, # attention_mask, # past_key_values, # inputs_embeds, # labels, # latents # ) = self.prepare_inputs_labels_for_multimodal( # input_ids, # position_ids, # attention_mask, # past_key_values, # labels, # gen_image, # und_image, # grid_thw, # i_s_pos, # image_sizes # ) # outputs = self.model( # input_ids=input_ids, # attention_mask=attention_mask, # position_ids=position_ids, # past_key_values=past_key_values, # inputs_embeds=inputs_embeds, # use_cache=use_cache, # output_attentions=output_attentions, # output_hidden_states=output_hidden_states, # return_dict=return_dict, # ) # hidden_states = outputs[0] # logits = self.lm_head(hidden_states) # logits = logits.float() # total_loss = None # if labels is not None: # # Shift so that tokens < n predict n # shift_logits = logits[..., :-1, :].contiguous() # shift_labels = labels[..., 1:].contiguous() # # Flatten the tokens # loss_fct = torch.nn.CrossEntropyLoss() # shift_logits = shift_logits.view(-1, self.config.vocab_size) # shift_labels = shift_labels.view(-1) # # Enable model parallelism # shift_labels = shift_labels.to(shift_logits.device) # loss = loss_fct(shift_logits, shift_labels) # # compute image loss # # target_img_embeds = torch.clone(inputs_embeds.detach())[:,1:,:] # get target image emb # img_loss_funct = torch.nn.MSELoss() # # img_hidden_states = self.get_model().down_projector(hidden_states[:,-self.get_n_query():,:]) # img_hidden_states = [] # for b in range(hidden_states.shape[0]): # img_hidden_states.append(hidden_states[b,i_s_pos[b]:i_s_pos[b]+64,:]) # img_hidden_states = torch.stack(img_hidden_states,dim=0) # img_hidden_states = self.get_model().down_projector(img_hidden_states) # # img_loss = 0.0 # if latents is None: # img_loss = img_loss_funct(img_hidden_states, torch.clone(img_hidden_states.detach())) # else: # bsz = latents.shape[0] # # device = latents.device # dtype = latents.dtype # noise = torch.randn_like(latents, device=latents.device) # u = torch.rand(size=(bsz,), device="cpu") # indices = (u * self.get_model().noise_scheduler.config.num_train_timesteps).long() # timesteps = self.get_model().noise_scheduler.timesteps[indices].to(device=latents.device) # sigmas = self.get_sigmas(timesteps, latents.device, n_dim=latents.ndim, dtype=dtype) # noisy_latents = (1.0 - sigmas) * latents + sigmas * noise # noise_pred = self.get_model().dit( # x=noisy_latents, # timestep=timesteps, # z_latents=self.mask_drop(img_hidden_states), # ) # target = noise - latents # img_loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean") # print(f"img loss {img_loss}") # total_loss = img_loss # return CausalLMOutputWithPast( # loss=total_loss, # logits=logits, # past_key_values=outputs.past_key_values, # hidden_states=outputs.hidden_states, # attentions=outputs.attentions, # ) # @torch.no_grad() # def generate( # self, # inputs: Optional[torch.Tensor] = None, # images: Optional[torch.Tensor] = None, # image_sizes: Optional[torch.Tensor] = None, # **kwargs, # ) -> Union[GenerateOutput, torch.LongTensor]: # position_ids = kwargs.pop("position_ids", None) # attention_mask = kwargs.pop("attention_mask", None) # if "inputs_embeds" in kwargs: # raise NotImplementedError("`inputs_embeds` is not supported") # if images is not None: # ( # inputs, # position_ids, # attention_mask, # _, # inputs_embeds, # img_indicator, # _ # ) = self.prepare_inputs_labels_for_understanding( # inputs, # position_ids, # attention_mask, # None, # None, # images, # image_sizes=image_sizes # ) # else: # inputs_embeds = self.get_model().embed_tokens(inputs) # return super().generate( # position_ids=position_ids, # attention_mask=attention_mask, # inputs_embeds=inputs_embeds, # **kwargs # ) @torch.no_grad() def generate_image( self, text: List[str], tokenizer: AutoTokenizer, pixel_values: Optional[torch.Tensor] = None, image_grid_thw: Optional[torch.Tensor] = None, max_var: Optional[float] = None, # placeholder: str = DEFAULT_IMG_PLACEHOLDER, ): scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", subfolder="scheduler") N_QUERY = self.get_n_query() inputs = tokenizer(text, padding="longest", return_tensors="pt") device = self.get_model().device attention_mask = inputs.attention_mask.to(device) input_ids = inputs.input_ids.to(device) # B x N input_ids = torch.cat([input_ids, torch.tensor([[151665]]).to(device)], dim=1) # breakpoint() text_embeds = self.get_model().embed_tokens(input_ids) latent_queries = self.get_model().latent_queries.repeat(text_embeds.shape[0], 1, 1) if pixel_values is not None: und_image_idx = (input_ids == UND_IMAGE_TOKEN_IDX) pixel_values = pixel_values.type(self.visual.dtype) und_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) text_embeds[und_image_idx] = und_image_embeds.to(text_embeds.device)[:und_image_idx.sum(), :] text_embeds = torch.cat([text_embeds, latent_queries], dim=1) attention_mask = torch.cat([attention_mask, torch.ones_like(latent_queries[:, :, 0])], dim=1) outputs = self.model( inputs_embeds=text_embeds, attention_mask=attention_mask, output_hidden_states=True, return_dict=True, ) hidden_states = outputs.hidden_states[-1][:,-N_QUERY:,:] img_hidden_states = hidden_states output_img = self.sample_images(img_hidden_states, scheduler) output_img = output_img.view(1, 1792, -1).permute(0,2,1).contiguous() return output_img def sample_images( self, img_hidden_states, scheduler, guidance_scale: float = 3.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, num_inference_steps: int = 30, num_images_per_prompt: int = 1, return_tensor=False, **kwargs, ): device = img_hidden_states.device dtype = img_hidden_states.dtype img_hidden_states_null = torch.zeros_like(img_hidden_states, device=device, dtype=dtype) img_hidden_states_input = torch.cat([img_hidden_states_null, img_hidden_states], 0) batch_size = img_hidden_states.shape[0] latent_size = self.get_model().dit.config.input_size latent_channels = self.get_model().dit.config.in_channels latents = randn_tensor( shape=(batch_size * num_images_per_prompt, latent_channels, latent_size, latent_size), generator=generator, device=device, dtype=dtype, ) # set step values sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) scheduler.set_timesteps(num_inference_steps, sigmas=sigmas) # Repeat z_latents and conditions for each image per prompt img_hidden_states_input = img_hidden_states_input.repeat_interleave(num_images_per_prompt, dim=0) for t in scheduler.timesteps: latent_model_input = latents.repeat(2, 1, 1, 1) if hasattr(scheduler, "scale_model_input"): latent_model_input = scheduler.scale_model_input(latent_model_input, t) # predict noise model_output noise_pred = self.get_model().dit( x=latent_model_input, timestep=t.unsqueeze(0).expand(latent_model_input.shape[0]).to(latent_model_input.device, torch.long), z_latents=img_hidden_states_input, ) # perform guidance noise_pred_uncond, noise_pred = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) # compute previous image: x_t -> x_t-1 latents = scheduler.step(noise_pred, t, latents).prev_sample # samples = self.decode_latents(latents, return_tensor=return_tensor) # breakpoint() return latents def decode_latents(self, latents, normalize=True, return_tensor=False): if isinstance(self.get_model().vae, AutoencoderKL): latents = latents / self.get_model().vae.config.scaling_factor if self.get_model().vae.config.shift_factor is not None: latents = latents + self.get_model().vae.config.shift_factor latents = latents.to(dtype=torch.float32) samples = self.get_model().vae.decode(latents).sample else: samples = self.get_model().vae.decode(latents) if normalize: samples = (samples / 2 + 0.5).clamp(0, 1) else: samples = samples.clamp(-1, 1) if return_tensor: return samples samples = samples.cpu().permute(0, 2, 3, 1).float().numpy() samples = numpy_to_pil(samples) return samples def prepare_and_encode_inputs( self, inputs: List[str | Image.Image], tokenizer: AutoTokenizer, do_classifier_free_guidance: bool = False, ): # pdb.set_trace() device = self.get_model().device dtype = self.get_model().dtype has_image, has_text = False, False text_prompt, image_prompt = "", [] img_processor = self.get_vision_tower().image_processor negative_prompt = {} for x in inputs: if isinstance(x, str): has_text = True text_prompt += x else: has_image = True text_prompt += DEFAULT_IMAGE_TOKEN image_prompt.append(img_processor.preprocess(x, return_tensors='pt')['pixel_values']) # pdb.set_trace() if len(image_prompt) == 0: image_prompt = None else: image_prompt = torch.cat(image_prompt) image_prompt = image_prompt.type(dtype).to(device) if has_image and not has_text: prompt = self.encode_images(image_prompt) # pdb.set_trace() if do_classifier_free_guidance: key = "[NULL_IMAGE]" if key not in negative_prompt: negative_image = torch.zeros_like(image_prompt) negative_prompt[key] = self.encode_images(negative_image) prompt = torch.cat([prompt, negative_prompt[key]], dim=0) else: prompt = self.generate_image(text=[text_prompt], image=image_prompt, tokenizer=tokenizer) if do_classifier_free_guidance: key = "" if key not in negative_prompt: negative_prompt[key] = self.generate_image(text=[""], tokenizer=tokenizer) prompt = torch.cat([prompt, negative_prompt[key]], dim=0) gen_pooling = self.get_gen_pooling() n_query = self.get_n_query() num_img, _, c = prompt.shape if 'pool2d' in gen_pooling and has_text and not 'early' in gen_pooling: stride = int(gen_pooling.split('_')[1]) sqrt_n = int(n_query**0.5) prompt = prompt.permute(0, 2, 1).reshape(num_img, -1, sqrt_n, sqrt_n) prompt = F.avg_pool2d(prompt, kernel_size=(stride, stride), stride=stride) prompt = prompt.reshape(num_img, c, -1).permute(0,2,1) return prompt def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs): images = kwargs.pop("images", None) image_sizes = kwargs.pop("image_sizes", None) inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs ) if images is not None: inputs['images'] = images if image_sizes is not None: inputs['image_sizes'] = image_sizes return inputs AutoConfig.register("blip3o_qwen", blip3oQwenConfig) AutoModelForCausalLM.register(blip3oQwenConfig, blip3oQwenForCausalLM)