update handling of init args

custom_st.py

@@ -22,6 +22,8 @@ class Transformer(nn.Module):
         min_pixels: int = 1 * 28 * 28,
         dimension: int = 2048,
         max_seq_length: Optional[int] = None,
+        model_args: Optional[Dict[str, Any]] = None,
+        processor_args: Optional[Dict[str, Any]] = None,
         cache_dir: Optional[str] = None,
         device: str = 'cuda:0',
         **kwargs,
@@ -34,6 +36,17 @@ class Transformer(nn.Module):
         self.min_pixels = min_pixels
         self.max_seq_length = max_seq_length
         
+        # Handle args
+        model_kwargs = model_args or {}
+        model_kwargs.update(kwargs)
+
+        processor_kwargs = processor_args or {}
+        processor_kwargs.update({
+            'min_pixels': min_pixels,
+            'max_pixels': max_pixels,
+            'cache_dir': cache_dir
+        })
+        
         # Initialize model
         try:
             self.model = Qwen2VLForConditionalGeneration.from_pretrained(
@@ -42,7 +55,7 @@ class Transformer(nn.Module):
                 torch_dtype=torch.bfloat16,
                 device_map=device,
                 cache_dir=cache_dir,
-                **kwargs
+                **model_kwargs
             ).eval()
         except (ImportError, ValueError) as e:
             print(f"Flash attention not available, falling back to default attention: {e}")
@@ -51,15 +64,13 @@ class Transformer(nn.Module):
                 torch_dtype=torch.bfloat16,
                 device_map=device,
                 cache_dir=cache_dir,
-                **kwargs
+                **model_kwargs
             ).eval()
 
         # Initialize processor
         self.processor = AutoProcessor.from_pretrained(
             processor_name_or_path or model_name_or_path,
-            min_pixels=min_pixels,
-            max_pixels=max_pixels,
-            cache_dir=cache_dir
+            **processor_kwargs
         )
 
         # Set padding sides