ruff fix

Signed-off-by: Amit Raj <quic_amitraj@quicinc.com>

Author: quic-amitraj

QEfficient/transformers/models/mllama/modeling_mllama.py

@@ -10,11 +10,9 @@
 import math
 from typing import List, Optional, Tuple, Union
 
-import requests
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
-from PIL import Image
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers.cache_utils import Cache, DynamicCache
@@ -1197,12 +1195,13 @@ def forward(
         return outputs
 
     def generate_input(self, processor, kv_offload):
-        
-        #vision_inputs
+        # vision_inputs
         vision_inputs = {
-            "pixel_values": torch.zeros((bs, max_num_images,max_image_tiles,num_channel, image_length, image_width ), dtype=torch.int64),
+            "pixel_values": torch.zeros(
+                (bs, max_num_images, max_image_tiles, num_channel, image_length, image_width), dtype=torch.int64
+            ),
             "aspect_ratio_ids": torch.ones((bs, max_num_images), dtype=torch.int64),
-            "aspect_ratio_mask": torch.ones((bs, max_num_images, max_image_tiles,1 ), dtype=torch.int64)
+            "aspect_ratio_mask": torch.ones((bs, max_num_images, max_image_tiles, 1), dtype=torch.int64),
         }
 
         vision_output_names = []
@@ -1220,19 +1219,19 @@ def generate_input(self, processor, kv_offload):
             },
         }
 
-        #lang_inputs
+        # lang_inputs
         lang_inputs = {
-            "input_ids": torch.zeros((bs,seq_len),dtype=torch.int64),
+            "input_ids": torch.zeros((bs, seq_len), dtype=torch.int64),
             "position_ids": torch.arange(seq_len, dtype=torch.int64).view(1, seq_len).repeat(bs, 1),
-            "cross_attention_mask": torch.ones((bs, max_image_tiles),dtype=torch.int64),
-            "attention_mask": torch.ones((bs,seq_len),dtype=torch.int64)
+            "cross_attention_mask": torch.ones((bs, max_image_tiles), dtype=torch.int64),
+            "attention_mask": torch.ones((bs, seq_len), dtype=torch.int64),
         }
 
         lang_inputs["position_ids"] = torch.where(
             lang_inputs.pop("attention_mask") == 1,
             torch.arange(lang_inputs["input_ids"].shape[1]).view(1, -1),
             -1,
-        )   
+        )
 
         ctx_len = Constants.CTX_LEN
         txt_cfg = self.mllama.config.get_text_config()
@@ -1245,7 +1244,6 @@ def generate_input(self, processor, kv_offload):
         num_patches = (vis_cfg.image_size // vis_cfg.patch_size) ** 2 + 1
         image_tokens_len = vis_cfg.max_num_tiles * num_patches
 
-
         lang_inputs["past_key_values"] = DynamicCache(num_hidden_layers)
         lang_inputs["past_key_values"].key_cache = [0] * num_hidden_layers
         lang_inputs["past_key_values"].value_cache = [0] * num_hidden_layers
@@ -1254,20 +1252,21 @@ def generate_input(self, processor, kv_offload):
             if i in cross_attention_layers:
                 idx = cross_attention_layers.index(i)
                 assert idx == ((i - 3) // 5), f"{i}, {(i - 3) // 5}"
-                lang_inputs["past_key_values"].key_cache[i] = torch.zeros(1, num_key_value_heads, image_tokens_len, head_dim)
+                lang_inputs["past_key_values"].key_cache[i] = torch.zeros(
+                    1, num_key_value_heads, image_tokens_len, head_dim
+                )
                 lang_inputs["past_key_values"].value_cache[i] = torch.zeros(
                     1, num_key_value_heads, image_tokens_len, head_dim
                 )
             else:
                 lang_inputs["past_key_values"].key_cache[i] = torch.zeros(1, num_key_value_heads, ctx_len, head_dim)
                 lang_inputs["past_key_values"].value_cache[i] = torch.zeros(1, num_key_value_heads, ctx_len, head_dim)
 
-        
         lang_output_names = [
             "logits",
             *[f"past_{kv}.{i}_RetainedState" for i in range(num_hidden_layers) for kv in ["key", "value"]],
         ]
-        
+
         lang_dynamic_axes = {
             "input_ids": {0: "batch_size", 1: "seq_len"},
             "position_ids": {0: "batch_size", 1: "seq_len"},
@@ -1286,10 +1285,10 @@ def generate_input(self, processor, kv_offload):
             else:
                 lang_dynamic_axes[f"past_key.{i}"] = {0: "batch_size", 2: "ctx_len"}
                 lang_dynamic_axes[f"past_value.{i}"] = {0: "batch_size", 2: "ctx_len"}
-        
+
         lang_inputs["past_key_values"] = lang_inputs["past_key_values"].to_legacy_cache()
         lang_inputs["position_ids"] = torch.full(lang_inputs["position_ids"].shape, ctx_len - 1)
-        
+
         inputs = []
         output_names = []
         dynamic_axes = []
@@ -1304,5 +1303,3 @@ def generate_input(self, processor, kv_offload):
             dynamic_axes.append({**vision_dynamic_axes, **lang_dynamic_axes})
 
         return inputs, output_names, dynamic_axes
-
-        

QEfficient/transformers/models/modeling_auto.py

@@ -14,11 +14,9 @@
 from typing import List, Optional, Union
 
 import numpy as np
-import requests
 import torch
 import torch.nn as nn
 import transformers
-from PIL import Image
 from transformers import (
     AutoModel,
     AutoModelForCausalLM,
@@ -34,7 +32,6 @@
 from QEfficient.base.onnx_transforms import FP16ClipTransform, SplitTensorsTransform
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.generation.text_generation_inference import get_compilation_dims
-from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.models.mllama.modeling_mllama import ModelWrapper, VisionEncoder
 from QEfficient.transformers.models.pytorch_transforms import CustomOpsTransform, KVCacheTransform, SpDTransform
 from QEfficient.transformers.quantizers.auto import QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING, with_replaced_quantizers
@@ -746,8 +743,7 @@ def export(
         self,
         export_dir: Optional[str] = None,
         **kwargs,
-    ) -> str:       
-
+    ) -> str:
         self.inputs, self.output_names, self.dynamic_axes = self.model.generate_input(self.processor)
         if self.kv_offload:
             self.vision_export_path = self.export_vision(export_dir)
@@ -757,12 +753,11 @@ def export(
             self._export(self.model, self.inputs[0], self.output_names[0], self.dynamic_axes[0], export_dir=export_dir)
 
     def export_vision(self, export_dir):
-        
-        self.vision_encoder_model=VisionEncoder(self.model)
+        self.vision_encoder_model = VisionEncoder(self.model)
 
-        vision_inputs=self.inputs[0]
-        vision_output_names=self.output_names[0]
-        vision_dynamic_axes=self.dynamic_axes[0]
+        vision_inputs = self.inputs[0]
+        vision_output_names = self.output_names[0]
+        vision_dynamic_axes = self.dynamic_axes[0]
 
         self.vision_onnx_path = self._export(
             self.vision_encoder_model,
@@ -775,20 +770,16 @@ def export_vision(self, export_dir):
         return self.vision_onnx_path
 
     def export_lang(self, export_dir):
-        self.lang_model= ModelWrapper(self.model)
+        self.lang_model = ModelWrapper(self.model)
 
-        lang_inputs=self.inputs[1]
-        lang_output_names=self.output_names[1]
-        lang_dynamic_axes=self.dynamic_axes[1]
+        lang_inputs = self.inputs[1]
+        lang_output_names = self.output_names[1]
+        lang_dynamic_axes = self.dynamic_axes[1]
 
         self.lang_onnx_path = self._export(
-            self.lang_model,
-            lang_inputs,
-            lang_output_names,
-            lang_dynamic_axes,
-            export_dir=export_dir
-            )
-        
+            self.lang_model, lang_inputs, lang_output_names, lang_dynamic_axes, export_dir=export_dir
+        )
+
         return self.lang_onnx_path
 
     def compile(