Upload A_MobileO_With_Edit/blip3o_fast_inference.py with huggingface_hub

A_MobileO_With_Edit/blip3o_fast_inference.py

@@ -0,0 +1,383 @@
+from typing import List, Optional
+import torch
+import torch.nn as nn
+from PIL import Image
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import Qwen2_5_VLConfig, Qwen2ForCausalLM, Qwen2Config, Qwen2Model
+from blip3o.constants import UND_IMAGE_TOKEN_IDX, DEFAULT_IMAGE_TOKEN
+from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import numpy_to_pil
+import numpy as np
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from tqdm import tqdm
+
+
+class blip3oFastConfig(Qwen2Config):
+    model_type = "blip3o_fast_inference"
+
+
+class blip3oFastModel(LlavaMetaModel, Qwen2Model):
+    config_class = blip3oFastConfig
+
+    def __init__(self, config: Qwen2_5_VLConfig):
+        super(blip3oFastModel, self).__init__(config)
+
+
+class blip3oFastForInferenceLM(Qwen2ForCausalLM, LlavaMetaForCausalLM):
+    config_class = blip3oFastConfig
+
+    def __init__(self, config):
+        super(blip3oFastForInferenceLM, self).__init__(config)
+        config.model_type = "blip3o_qwen_inference"
+
+        self.model = blip3oFastModel(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 visual(self, pixel_values: torch.Tensor, grid_thw: Optional[torch.Tensor] = None) -> torch.Tensor:
+        image_features = self.get_model().get_vision_tower()(pixel_values)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+    
+    @torch.no_grad()
+    def generate_image(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_grid_thw: Optional[torch.Tensor] = None,
+        with_cfg: bool = True,
+        max_var: Optional[float] = None,
+        num_inference_steps: int = 20,
+    ):  
+        text_embeds = self.get_model().embed_tokens(input_ids)
+
+
+        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(), :]
+
+        outputs = self.model(
+            inputs_embeds=text_embeds,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        img_hidden_states = outputs.hidden_states
+        output_img = self.sample_images(img_hidden_states, attention_mask, with_cfg, num_inference_steps=num_inference_steps)
+        return output_img
+    def sample_images(
+            self,
+            pred_latents,  # Tuple/list of hidden states from all layers
+            attention_mask,
+            with_cfg: bool = True,
+            guidance_scale: float = 1.2,
+            num_inference_steps: int = 20,
+            num_images_per_prompt: int = 1,
+            return_tensor=False,
+            with_tqdm: bool = True,
+            **kwargs,
+    ):
+        # Get device and dtype from first element of tuple
+        device = pred_latents[0].device
+        dtype = pred_latents[0].dtype
+
+        # ✅ Store original batch size BEFORE any CFG operations
+        batch_size = pred_latents[0].shape[0]
+
+        latent_size = self.get_model().dit.config.sample_size
+        latent_channels = self.get_model().dit.config.in_channels
+
+        # ================================================================
+        # 1. CFG Preparation (same as sample_images_with_sam3)
+        # ================================================================
+        if with_cfg:
+            pred_latents_cfg = tuple(
+                torch.cat([torch.zeros_like(layer), layer], dim=0)
+                for layer in pred_latents
+            )
+        else:
+            pred_latents_cfg = pred_latents
+
+        # ================================================================
+        # 2. Process through connector
+        # ================================================================
+        encoder_hidden_states = self.model.diffusion_connector(pred_latents_cfg)
+        # Shape: [B, N, hidden_dim] or [2*B, N, hidden_dim] if with_cfg
+
+        # ================================================================
+        # 3. Initialize Latents (use original batch_size, not doubled)
+        # ================================================================
+        latents = randn_tensor(
+            shape=(batch_size * num_images_per_prompt, latent_channels, latent_size, latent_size),
+            generator=None,
+            device=device,
+            dtype=dtype,
+        )
+
+        # ================================================================
+        # 4. Denoising Loop
+        # ================================================================
+        self.model.noise_scheduler.set_timesteps(num_inference_steps)
+
+        iterator = tqdm(self.model.noise_scheduler.timesteps,
+                        desc="Sampling") if with_tqdm else self.model.noise_scheduler.timesteps
+
+        for t in iterator:
+            # Prepare model input
+            if with_cfg:
+                latent_model_input = torch.cat([latents] * 2)
+            else:
+                latent_model_input = latents
+
+            latent_model_input = latent_model_input.to(dtype)
+
+            # Scale model input if needed
+            if hasattr(self.model.noise_scheduler, "scale_model_input"):
+                latent_model_input = self.model.noise_scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+            # DiT forward
+            noise_pred = self.model.dit(
+                hidden_states=latent_model_input,
+                encoder_hidden_states=encoder_hidden_states,
+                timestep=t.unsqueeze(0).expand(latent_model_input.shape[0]).to(device),
+                encoder_attention_mask=None
+            ).sample
+
+            # Apply classifier-free guidance
+            if with_cfg:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # Perform denoising step
+            latents = self.model.noise_scheduler.step(noise_pred, t, latents).prev_sample
+
+        # ================================================================
+        # 5. Decode
+        # ================================================================
+        samples = self.decode_latents(latents.to(self.model.vae.dtype), return_tensor=return_tensor)
+        return samples
+    def sample_images_(
+            self,
+            pred_latents,  # Tuple/list of hidden states from all layers
+            attention_mask,
+            with_cfg: bool = False,
+            guidance_scale: float = 3.0,
+            num_inference_steps: int = 30,
+            num_images_per_prompt: int = 1,
+            return_tensor=False,
+            **kwargs,
+    ):
+        # Get device and dtype from first element of tuple
+        device = pred_latents[0].device
+        dtype = pred_latents[0].dtype
+    
+        # Compute conditioning ONCE before the loop (not inside!)
+        encoder_hidden_states = self.model.diffusion_connector(pred_latents)
+        # Shape: [B, N, 2304]
+    
+        # Handle classifier-free guidance
+        if with_cfg:
+            # Create null conditioning (zeros) for CFG
+            encoder_hidden_states_null = torch.zeros_like(encoder_hidden_states)
+            # Concatenate: [null_batch, text_batch]
+            encoder_hidden_states = torch.cat([encoder_hidden_states_null, encoder_hidden_states], dim=0)
+    
+        # Get batch size from processed encoder hidden states
+        batch_size = encoder_hidden_states.shape[0]
+        latent_size = self.get_model().dit.config.sample_size
+        latent_channels = self.get_model().dit.config.in_channels
+    
+        # Initialize random noise latents
+        latents = randn_tensor(
+            shape=(batch_size * num_images_per_prompt, latent_channels, latent_size, latent_size),
+            generator=None,
+            device=device,
+            dtype=dtype,
+        )
+
+    def sample_images_original(
+            self,
+            pred_latents,  # Tuple/list of hidden states from all layers
+            attention_mask,
+            with_cfg: bool = False,
+            guidance_scale: float = 1.2,
+            num_inference_steps: int = 30,
+            num_images_per_prompt: int = 1,
+            return_tensor=False,
+            **kwargs,
+    ):
+        # Get device and dtype from first element of tuple
+        device = pred_latents[0].device
+        dtype = pred_latents[0].dtype  # ✅ Store dtype here
+    
+        encoder_hidden_states = self.model.diffusion_connector(pred_latents)
+        # Shape: [B, N, 2304]
+    
+        # Handle classifier-free guidance
+        if with_cfg:
+            # Create null conditioning (zeros) for CFG
+            encoder_hidden_states_null = torch.zeros_like(encoder_hidden_states)
+            # Concatenate: [null_batch, text_batch]
+            encoder_hidden_states = torch.cat([encoder_hidden_states_null, encoder_hidden_states], dim=0)
+    
+        # Get batch size from processed encoder hidden states
+        batch_size = encoder_hidden_states.shape[0]
+        latent_size = self.get_model().dit.config.sample_size
+        latent_channels = self.get_model().dit.config.in_channels
+    
+        # Initialize random noise latents
+        latents = randn_tensor(
+            shape=(batch_size * num_images_per_prompt, latent_channels, latent_size, latent_size),
+            generator=None,
+            device=device,
+            dtype=dtype,
+        )
+    
+        # Set up timesteps for denoising
+        self.model.noise_scheduler.set_timesteps(num_inference_steps)
+    
+        # Denoising loop
+        for t in tqdm(self.model.noise_scheduler.timesteps, desc="Sampling images"):
+            # Prepare model input
+            if with_cfg:
+                # Duplicate latents for CFG (unconditional + conditional)
+                latent_model_input = torch.cat([latents] * 2)
+            else:
+                latent_model_input = latents
+    
+            latent_model_input = latent_model_input.to(dtype)
+    
+            # Scale model input if needed (for some schedulers)
+            if hasattr(self.model.noise_scheduler, "scale_model_input"):
+                latent_model_input = self.model.noise_scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+    
+            # Predict noise
+            noise_pred = self.model.dit(
+                hidden_states=latent_model_input,
+                encoder_hidden_states=encoder_hidden_states,
+                timestep=t.unsqueeze(0).expand(latent_model_input.shape[0]).to(device),
+                encoder_attention_mask=None
+            ).sample
+    
+            # Apply classifier-free guidance
+            if with_cfg:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+    
+            # Perform denoising step
+            latents = self.model.noise_scheduler.step(noise_pred, t, latents).prev_sample
+    
+        # Decode latents to images
+        samples = self.decode_latents(latents.to(self.model.vae.dtype), return_tensor=return_tensor)
+        return samples
+
+
+    @torch.no_grad()
+    def decode_latents(self, latents, normalize=True, return_tensor=False):
+        if self.model.vae is not None:
+            latents = latents / self.model.vae.config.scaling_factor
+            if "shift_factor" in self.model.vae.config and self.model.vae.config.shift_factor is not None:
+                latents = latents + self.model.vae.config.shift_factor
+            samples = self.model.vae.decode(latents).sample
+        else:
+            samples = 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,
+    ): 
+        print("="*20, "prepare_and_encode_inputs", "="*20)
+        # 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):
+        print("="*20, "prepare_inputs_for_generation", "="*20)
+        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_fast_inference", blip3oFastConfig)
+AutoModelForCausalLM.register(blip3oFastConfig, blip3oFastForInferenceLM)