Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test_image.png filter=lfs diff=lfs merge=lfs -text

LICENSE.md

@@ -0,0 +1,94 @@
+# License for SafeLLaVA
+
+The SafeLLaVA project is governed by a **hybrid license model**. This license file defines the distinct licensing policies that apply to the two main components of this project.
+
+---
+
+## 1. Definition of Work
+
+### Model Name
+**SafeLLaVA**
+
+### Reference Publication
+This model (SafeLLaVA) is the official model presented in the academic paper:
+
+> **"HoliSafe: Holistic Safety Benchmarking and Modeling for Vision-Language Model"**  
+> [https://arxiv.org/abs/2506.04704](https://arxiv.org/abs/2506.04704)
+
+### Base Model
+This model is a **Derivative Work** based on the LLaVA-v1.5 model:
+- [https://github.com/haotian-liu/LLaVA](https://github.com/haotian-liu/LLaVA)
+
+### Modifications by ETRI
+This work integrates an independently developed **Visual Guard Module (VGM)** to classify harmful image inputs and generate safe text responses. All modifications and additions are the work of ETRI.
+
+---
+
+## 2. License Summary
+
+| Component | License |
+|-----------|---------|
+| Independently Developed Code (e.g., VGM) | Apache License 2.0 |
+| LLaVA-Based Components and Entire Model | LLaVA-v1.5 Usage and License Notices |
+
+---
+
+## Part 1: Apache License 2.0 (For Independently Developed Code)
+
+All original source code and components developed independently by **Electronics and Telecommunications Research Institute (ETRI)** (hereinafter "Copyright Holder"), including the **Visual Guard Module (VGM)** contained in this project, are subject to the Apache License, Version 2.0 (the "License").
+
+You may not use this file except in compliance with the License. You may obtain a copy of the License at:
+
+[http://www.apache.org/licenses/LICENSE-2.0](http://www.apache.org/licenses/LICENSE-2.0)
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
+```
+Copyright 2025 Electronics and Telecommunications Research Institute (ETRI)
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+```
+
+---
+
+## Part 2: LLaVA-v1.5 Usage and License Notices (For the LLaVA-Based Derivative Work)
+
+This SafeLLaVA model is a **Derivative Work** based on the LLaVA-v1.5 model.
+
+Therefore, the use, reproduction, modification, and distribution of the entire SafeLLaVA model (including LLaVA-based components and weights) are subject to the original **"Usage and License Notices"** from the LLaVA-v1.5 repository.
+
+Any user of SafeLLaVA must agree to and comply with these upstream notices, which are reproduced below:
+
+### Usage and License Notices (from [LLaVA-v1.5](https://github.com/haotian-liu/LLaVA?tab=readme-ov-file))
+
+> This project utilizes certain datasets and checkpoints that are subject to their respective original licenses. Users must comply with all terms and conditions of these original licenses, including but not limited to the **OpenAI Terms of Use** for the dataset and the specific licenses for base language models for checkpoints trained using the dataset (e.g. **Llama community license** for LLaMA-2 and **Vicuna-v1.5**). This project does not impose any additional constraints beyond those stipulated in the original licenses. Furthermore, users are reminded to ensure that their use of the dataset and checkpoints is in compliance with all applicable laws and regulations.
+
+### Implications for Users
+
+This means that any user of SafeLLaVA is responsible for independently verifying and complying with all applicable upstream licenses, which may include (but are not limited to):
+
+- **The OpenAI Terms of Use**  
+  [https://openai.com/policies/terms-of-use](https://openai.com/policies/terms-of-use)
+
+- **The Llama 2 Community License**  
+  [https://ai.meta.com/llama/license/](https://ai.meta.com/llama/license/)
+
+---
+
+## Part 3: Attribution and Contact
+
+This SafeLLaVA model was developed by the **Electronics and Telecommunications Research Institute (ETRI)** in the Republic of Korea.
+
+For any questions regarding the SafeLLaVA model or its licensing, please contact:
+
+**Youngwan Lee**  
+Email: [[email protected]](mailto:[email protected])

README.md

@@ -0,0 +1,323 @@
+---
+base_model: llava-v1.5
+tags:
+- vision
+- multimodal
+- safety
+- content-moderation
+- llava
+- image-classification
+- vision-language
+language:
+- en
+pipeline_tag: image-text-to-text
+library_name: transformers
+---
+
+# SafeLLaVA-13B: Vision-Language Model with Visual Guard Module
+
+[**🌐 Website**](https://youngwanlee.github.io/holisafe) | [**📑 Paper**](https://www.arxiv.org/pdf/2506.04704)
+
+
+
+<div align="center">
+  <img src="https://dl.dropbox.com/scl/fi/soi772p6sig2tx16f092o/arch.jpg?rlkey=uj4ver4pp889oowigqld502hc&dl=1" width="1024px" />
+</div>
+
+SafeLLaVA-13B is a safe multimodal large language model that extends [LLaVA-v1.5](https://github.com/haotian-liu/LLaVA) with built-in image safety classification capabilities. It can simultaneously generate text responses to visual questions while classifying potentially unsafe image content across 20 safety categories.
+
+## Model Description
+
+- **Base Model**: LLaVA-v1.5-13B
+- **Architecture**: Vision-language model with Visual Guard Module (VGM)
+- **Training Data**: HoliSafe train set
+- **Training Method**: LoRA fine-tuning
+- **Parameters**: 13B (base) + 8.43M (VGM)
+- **Safety Categories**: 20 categories based on HoliSafe taxonomy
+
+## Key Features
+
+1. **Multimodal Understanding**: Processes images and text for comprehensive visual understanding
+2. **Safety Classification**: Identifies unsafe content in images across 20 categories
+3. **Non-invasive Architecture**: Maintains full LLaVA-v1.5-13B capabilities while adding safety features
+4. **End-to-end Training**: VGM is jointly trained with the vision-language model
+
+The model classifies images into the following 20 safety categories:
+
+| Category ID | Category Name |
+|------------|---------------|
+| 0 | Safe |
+| 1 | Gender discrimination |
+| 2 | Race discrimination |
+| 3 | Religion discrimination |
+| 4 | Harassment |
+| 5 | Disability discrimination |
+| 6 | Drug Related Hazards |
+| 7 | Property crime |
+| 8 | Facial data exposure |
+| 9 | Identity data exposure |
+| 10 | Physical self-injury |
+| 11 | Suicide |
+| 12 | Animal abuse |
+| 13 | Obscene gestures |
+| 14 | Physical altercation |
+| 15 | Terrorism |
+| 16 | Weapon-related violence |
+| 17 | Sexual content |
+| 18 | Financial advice |
+| 19 | Medical advice |
+
+## 🚀 Quick Start
+
+### Installation
+
+```bash
+pip install torch transformers pillow accelerate requests
+```
+
+### Complete Example - Copy & Paste Ready
+
+```python
+import requests
+import torch
+import sys
+from pathlib import Path
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+from huggingface_hub import snapshot_download, hf_hub_download
+
+# Model path
+model_path = "etri-vilab/SafeLLaVA-13B"
+
+# Download model and add safellava package to path
+model_cache_path = Path(snapshot_download(repo_id=model_path))
+sys.path.insert(0, str(model_cache_path))
+
+# Import safellava utilities
+from safellava.mm_utils import tokenizer_image_token
+from safellava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from safellava.conversation import conv_templates
+
+# Load model and tokenizer
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True
+)
+model = model.to('cuda:0')
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+
+# Load and move vision tower to GPU
+vision_tower = model.get_vision_tower()
+if not vision_tower.is_loaded:
+    vision_tower.load_model()
+vision_tower = vision_tower.to('cuda:0')
+
+print("✅ Model loaded successfully!")
+
+# Helper function to load image from URL or local path
+def load_image(image_file):
+    if image_file.startswith('http'):
+        from io import BytesIO
+        response = requests.get(image_file, timeout=30)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content)).convert('RGB')
+    else:
+        return Image.open(image_file).convert('RGB')
+
+# Download and load the test image from HuggingFace Hub
+# (The image is included in the model repository)
+test_image_path = hf_hub_download(repo_id=model_path, filename="test_image.png", repo_type="model")
+image = load_image(test_image_path)
+
+# You can also use your own image:
+# image = load_image("path/to/your/image.jpg")
+# Or load from URL:
+# image = load_image("https://example.com/image.jpg")
+
+# Preprocess image
+image_processor = vision_tower.image_processor
+image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values']
+image_tensor = image_tensor.to('cuda:0', dtype=torch.float16)
+
+# Prepare conversation prompt
+conv = conv_templates["llava_v1"].copy()
+question = "How to get this?"
+conv.append_message(conv.roles[0], DEFAULT_IMAGE_TOKEN + "\n" + question)
+conv.append_message(conv.roles[1], None)
+prompt = conv.get_prompt()
+
+input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')
+input_ids = input_ids.unsqueeze(0).to('cuda:0')
+
+# Run safety classification
+with torch.inference_mode():
+    outputs = model(
+        input_ids=input_ids,
+        images=image_tensor,
+        do_safety=True,
+        output_hidden_states=True,
+        return_dict=True
+    )
+
+# Get safety classification results
+if outputs.img_safety_probs is not None:
+    probs = outputs.img_safety_probs[0]
+    pred_class = probs.argmax().item()
+    pred_prob = probs[pred_class].item()
+
+    safety_categories = model.config.safety_categories
+    print(f"\n✅ Safety Classification Result:")
+    print(f"   Predicted: {safety_categories[pred_class]} ({pred_prob:.4f})")
+
+    # Show top 3 predictions
+    print("\n   Top 3 predictions:")
+    top_k = min(3, len(probs))
+    top_probs, top_indices = torch.topk(probs, k=top_k)
+    for i, (idx, prob) in enumerate(zip(top_indices, top_probs)):
+        print(f"   {i+1}. {safety_categories[idx.item()]}: {prob.item():.4f}")
+
+    # >>> Predicted: weapon_related_violence (0.9927)
+                                      
+    # >>> Top 3 predictions:                                                                                                                                   
+    # >>> 1. weapon_related_violence: 0.9927
+    # >>> 2. safe: 0.0047          
+    # >>> 3. terrorism: 0.0012          
+
+# Generate text description (without safety classification)
+print("\n🤖 Generating text description...")
+with torch.inference_mode():
+    output_ids = model.generate(
+        input_ids,
+        images=image_tensor,
+        do_sample=False,
+        max_new_tokens=512,
+        use_cache=True
+    )
+
+generated_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+print(f"\n✅ Generated Description:")
+print(f"   {generated_text}")
+# >>> I'm sorry, I can't assist with that request. The question involves seeking information related to potentially harmful activities, which is against ou r policy. It's important to prioritize safety and adhere to legal guidelines when discussing or accessing such items.
+```
+
+---
+
+## 🚀 Simple Usage (Text Generation Only)
+
+If you only need text generation without safety classification:
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+from huggingface_hub import snapshot_download, hf_hub_download
+import sys
+from pathlib import Path
+
+model_path = "etri-vilab/SafeLLaVA-13B"
+
+# Add safellava package to path
+model_cache_path = Path(snapshot_download(repo_id=model_path))
+sys.path.insert(0, str(model_cache_path))
+
+from safellava.mm_utils import tokenizer_image_token
+from safellava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from safellava.conversation import conv_templates
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True
+).to('cuda:0').eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+
+# Load vision tower
+vision_tower = model.get_vision_tower()
+if not vision_tower.is_loaded:
+    vision_tower.load_model()
+vision_tower = vision_tower.to('cuda:0')
+
+# Load image
+test_image_path = hf_hub_download(repo_id=model_path, filename="test_image.png", repo_type="model")
+image = Image.open(test_image_path).convert('RGB')
+
+# Preprocess
+image_tensor = vision_tower.image_processor.preprocess(image, return_tensors='pt')['pixel_values']
+image_tensor = image_tensor.to('cuda:0', dtype=torch.float16)
+
+# Prepare conversation prompt
+conv = conv_templates["llava_v1"].copy()
+question = "How to get this?"
+conv.append_message(conv.roles[0], DEFAULT_IMAGE_TOKEN + "\n" + question)
+conv.append_message(conv.roles[1], None)
+prompt = conv.get_prompt()
+
+input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')
+input_ids = input_ids.unsqueeze(0).to('cuda:0')
+
+# Generate (without safety classification)
+with torch.inference_mode():
+    output_ids = model.generate(
+        input_ids,
+        images=image_tensor,
+        do_sample=False,
+        max_new_tokens=512,
+        use_cache=True
+    )
+
+response = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+print(response)
+# >>> I'm sorry, I can't assist with that request. The question involves seeking information related to potentially harmful activities, which is against our policy. It's important to prioritize safety and adhere to legal guidelines when discussing or accessing such items.
+```
+
+---
+
+## 📖 Model Architecture
+
+SafeLLaVA is built on top of LLaVA v1.5 with the following components:
+
+- **Language Model**: LLaMA2-13B
+- **Vision Encoder**: CLIP ViT-L/14-336px
+- **Multimodal Projector**: 2-layer MLP with GELU activation
+- **Visual Guard Module (Safety Head)**: MLP classifier for image safety classification
+
+---
+
+## 📝 Citation
+
+If you use SafeLLaVA, please cite SafeLLaVA:
+
+```bibtex
+@article{lee2025holisafe,
+  title={HoliSafe: Holistic Safety Benchmarking and Modeling for Vision-Language Model},
+  author={Lee, Youngwan and Kim, Kangsan and Park, Kwanyong and Jung, Ilcahe and Jang, Soojin and Lee, Seanie and Lee, Yong-Ju and Hwang, Sung Ju},
+  journal={arXiv preprint arXiv:2506.04704},
+  year={2025},
+  url={https://arxiv.org/abs/2506.04704},
+  archivePrefix={arXiv},
+  eprint={2506.04704},
+  primaryClass={cs.AI},
+}
+```
+
+---
+
+## 📄 License
+
+See [LICENSE](LICENSE.md) for details.
+
+---
+
+## 🙏 Acknowledgments
+
+- Built on [LLaVA-v1.5](https://github.com/haotian-liu/LLaVA)
+
+This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. RS-2022-00187238, Development of Large Korean Language Model Technology for Efficient Pre-training, 45%), (No. 2022-0-00871, Development of AI Autonomy and Knowledge Enhancement for AI Agent Collaboration, 45%) and (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST), 45%).

model-00001-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30b19ff7d43c0f59522f07d280755e0cd47f36007cb5f3d63ce01d7600fcb7e7
+size 4978265728

model-00002-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f067f75a031295b4ee0e14da3990c31fa2c5a0654b02154945ce31a2a333385c
+size 4970422160

model-00003-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d2cf9d000b4d06888153032381e7607641e9e76317b81d100cc042618bfb574
+size 4970422184

model-00004-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f04bd9e776835ab2724793be821ce0ed764c5447631216df4bc8dc19076af3d5
+size 4933701432

model-00005-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5da5854596848ba5a512f193bd331185eb3efa911179c091ed81153a6567682d
+size 4933722144

model-00006-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff21d3d163affb34c2ad4d7021c93e014f301eb138c0f30493bc171ce0c5c9f3
+size 1941570592

modeling_safellava.py

@@ -0,0 +1,21 @@
+"""
+SafeLLaVA Model for HuggingFace Hub
+
+This model is based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Licensed under Apache License 2.0
+
+SafeLLaVA adds image safety classification capabilities to LLaVA.
+"""
+
+# Re-export classes from safellava package for HuggingFace auto_map
+from safellava.model.language_model.safe_llava_llama import (
+    SafetyConfig,
+    SafeLlavaLlamaForCausalLM,
+    SafetyCausalLMOutputWithPast,
+)
+
+__all__ = [
+    "SafetyConfig",
+    "SafeLlavaLlamaForCausalLM",
+    "SafetyCausalLMOutputWithPast",
+]

safellava/__init__.py

@@ -0,0 +1,11 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from .model.language_model.llava_llama import LlavaLlamaForCausalLM
+from .model.language_model.safe_llava_llama import SafeLlavaLlamaForCausalLM
+
+__all__ = ['LlavaLlamaForCausalLM', 'SafeLlavaLlamaForCausalLM']

safellava/constants.py

@@ -0,0 +1,20 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+IMAGE_PLACEHOLDER = "<image-placeholder>"

safellava/conversation.py

@@ -0,0 +1,96 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+"""
+Conversation prompts for SafeLLaVA.
+
+This is a simplified version containing only the llava_v1 conversation template.
+"""
+
+import dataclasses
+from enum import auto, Enum
+from typing import List
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    TWO = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.TWO
+    sep: str = " "
+    sep2: str = "</s>"
+    version: str = "v1"
+
+    def get_prompt(self):
+        """Generate the full prompt from conversation history."""
+        messages = self.messages
+
+        # Handle image token in first message
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role, message):
+        """Append a message to conversation history."""
+        self.messages.append([role, message])
+
+    def copy(self):
+        """Create a copy of this conversation."""
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+        )
+
+
+# LLaVA v1 conversation template
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+# Available conversation templates
+conv_templates = {
+    "llava_v1": conv_llava_v1,
+    "default": conv_llava_v1,
+}

safellava/mm_utils.py

@@ -0,0 +1,254 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from PIL import Image
+from io import BytesIO
+import base64
+import torch
+import math
+import ast
+
+from transformers import StoppingCriteria
+from safellava.constants import IMAGE_TOKEN_INDEX
+
+
+def select_best_resolution(original_size, possible_resolutions):
+    """
+    Selects the best resolution from a list of possible resolutions based on the original size.
+
+    Args:
+        original_size (tuple): The original size of the image in the format (width, height).
+        possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
+
+    Returns:
+        tuple: The best fit resolution in the format (width, height).
+    """
+    original_width, original_height = original_size
+    best_fit = None
+    max_effective_resolution = 0
+    min_wasted_resolution = float('inf')
+
+    for width, height in possible_resolutions:
+        scale = min(width / original_width, height / original_height)
+        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)
+        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
+        wasted_resolution = (width * height) - effective_resolution
+
+        if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
+            max_effective_resolution = effective_resolution
+            min_wasted_resolution = wasted_resolution
+            best_fit = (width, height)
+
+    return best_fit
+
+
+def resize_and_pad_image(image, target_resolution):
+    """
+    Resize and pad an image to a target resolution while maintaining aspect ratio.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        target_resolution (tuple): The target resolution (width, height) of the image.
+
+    Returns:
+        PIL.Image.Image: The resized and padded image.
+    """
+    original_width, original_height = image.size
+    target_width, target_height = target_resolution
+
+    scale_w = target_width / original_width
+    scale_h = target_height / original_height
+
+    if scale_w < scale_h:
+        new_width = target_width
+        new_height = min(math.ceil(original_height * scale_w), target_height)
+    else:
+        new_height = target_height
+        new_width = min(math.ceil(original_width * scale_h), target_width)
+
+    resized_image = image.resize((new_width, new_height))
+
+    new_image = Image.new('RGB', (target_width, target_height), (0, 0, 0))
+    paste_x = (target_width - new_width) // 2
+    paste_y = (target_height - new_height) // 2
+    new_image.paste(resized_image, (paste_x, paste_y))
+
+    return new_image
+
+
+def divide_to_patches(image, patch_size):
+    """
+    Divides an image into patches of a specified size.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        patch_size (int): The size of each patch.
+
+    Returns:
+        list: A list of PIL.Image.Image objects representing the patches.
+    """
+    patches = []
+    width, height = image.size
+    for i in range(0, height, patch_size):
+        for j in range(0, width, patch_size):
+            box = (j, i, j + patch_size, i + patch_size)
+            patch = image.crop(box)
+            patches.append(patch)
+
+    return patches
+
+
+def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
+    """
+    Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
+
+    Args:
+        image_size (tuple): The size of the input image in the format (width, height).
+        grid_pinpoints (str): A string representation of a list of possible resolutions.
+        patch_size (int): The size of each image patch.
+
+    Returns:
+        tuple: The shape of the image patch grid in the format (width, height).
+    """
+    if type(grid_pinpoints) is list:
+        possible_resolutions = grid_pinpoints
+    else:
+        possible_resolutions = ast.literal_eval(grid_pinpoints)
+    width, height = select_best_resolution(image_size, possible_resolutions)
+    return width // patch_size, height // patch_size
+
+
+def process_anyres_image(image, processor, grid_pinpoints):
+    """
+    Process an image with variable resolutions.
+
+    Args:
+        image (PIL.Image.Image): The input image to be processed.
+        processor: The image processor object.
+        grid_pinpoints (str): A string representation of a list of possible resolutions.
+
+    Returns:
+        torch.Tensor: A tensor containing the processed image patches.
+    """
+    if type(grid_pinpoints) is list:
+        possible_resolutions = grid_pinpoints
+    else:
+        possible_resolutions = ast.literal_eval(grid_pinpoints)
+    best_resolution = select_best_resolution(image.size, possible_resolutions)
+    image_padded = resize_and_pad_image(image, best_resolution)
+
+    patches = divide_to_patches(image_padded, processor.crop_size['height'])
+
+    image_original_resize = image.resize((processor.size['shortest_edge'], processor.size['shortest_edge']))
+
+    image_patches = [image_original_resize] + patches
+    image_patches = [processor.preprocess(image_patch, return_tensors='pt')['pixel_values'][0]
+                     for image_patch in image_patches]
+    return torch.stack(image_patches, dim=0)
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg):
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    new_images = []
+    if image_aspect_ratio == 'pad':
+        for image in images:
+            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            new_images.append(image)
+    elif image_aspect_ratio == "anyres":
+        for image in images:
+            image = process_anyres_image(image, image_processor, model_cfg.image_grid_pinpoints)
+            new_images.append(image)
+    else:
+        return image_processor(images, return_tensors='pt')['pixel_values']
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    offset = 0
+    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+
+    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+        input_ids.extend(x[offset:])
+
+    if return_tensors is not None:
+        if return_tensors == 'pt':
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f'Unsupported tensor type: {return_tensors}')
+    return input_ids
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith('checkpoint-'):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def call_for_batch(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            truncated_output_ids = output_ids[0, -keyword_id.shape[0]:]
+            if torch.equal(truncated_output_ids, keyword_id):
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        outputs = []
+        for i in range(output_ids.shape[0]):
+            outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+        return all(outputs)

safellava/model/__init__.py

@@ -0,0 +1,11 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from .language_model.llava_llama import LlavaLlamaForCausalLM, LlavaConfig
+from .language_model.safe_llava_llama import SafeLlavaLlamaForCausalLM
+
+__all__ = ['LlavaLlamaForCausalLM', 'SafeLlavaLlamaForCausalLM', 'LlavaConfig']

safellava/model/language_model/__init__.py

@@ -0,0 +1,11 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from .llava_llama import LlavaLlamaForCausalLM, LlavaConfig
+from .safe_llava_llama import SafeLlavaLlamaForCausalLM, SafetyConfig
+
+__all__ = ['LlavaLlamaForCausalLM', 'LlavaConfig', 'SafeLlavaLlamaForCausalLM', 'SafetyConfig']

safellava/model/language_model/llava_llama.py

@@ -0,0 +1,166 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation.utils import GenerateOutput
+
+from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+
+
+class LlavaConfig(LlamaConfig):
+    model_type = "llava_llama"
+
+
+class LlavaLlamaModel(LlavaMetaModel, LlamaModel):
+    config_class = LlavaConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(LlavaLlamaModel, self).__init__(config)
+
+
+class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
+    config_class = LlavaConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = LlavaLlamaModel(config)
+        self.pretraining_tp = config.pretraining_tp
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        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,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[int] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[List[List[int]]] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images,
+                image_sizes
+            )
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+    @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,
+                _
+            ) = self.prepare_inputs_labels_for_multimodal(
+                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
+        )
+
+    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("llava_llama", LlavaConfig)
+AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)

safellava/model/language_model/safe_llava_llama.py

@@ -0,0 +1,426 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from typing import List, Optional, Tuple, Union, Dict
+
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from safellava.model.language_model.llava_llama import (
+    LlavaConfig, LlavaLlamaModel, LlavaLlamaForCausalLM
+)
+from safellava.constants import IMAGE_TOKEN_INDEX
+
+from dataclasses import dataclass
+
+import logging
+from safellava.utils import setup_simple_logging
+
+setup_simple_logging()
+
+
+@dataclass
+class SafetyCausalLMOutputWithPast(CausalLMOutputWithPast):
+    """
+    Base class for causal language model (or autoregressive) outputs with safety predictions.
+    """
+    img_safety_logits: Optional[torch.FloatTensor] = None
+    img_safety_probs: Optional[torch.FloatTensor] = None
+    txt_safety_logits: Optional[torch.FloatTensor] = None
+    txt_safety_probs: Optional[torch.FloatTensor] = None
+    total_safety_logits: Optional[torch.FloatTensor] = None
+    total_safety_probs: Optional[torch.FloatTensor] = None
+
+
+class SafetyMLP(nn.Module):
+    """
+    Safety classification head implemented as Multi-layer Perceptron.
+    """
+
+    def __init__(self, input_size: int, hidden_size: int, output_size: int,
+                 safety_num_hidden_layers: int = 1):
+        super().__init__()
+
+        layers = []
+
+        layers.append(nn.Linear(input_size, hidden_size))
+        layers.append(nn.GELU())
+
+        for _ in range(safety_num_hidden_layers - 1):
+            layers.append(nn.Linear(hidden_size, hidden_size))
+            layers.append(nn.GELU())
+
+        layers.append(nn.Linear(hidden_size, output_size))
+
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class SafetyConfig(LlavaConfig):
+    """Safety-aware configuration for pooling version """
+    model_type = "safe_llava_llama"
+
+    def __init__(
+        self,
+        safety_categories=None,
+        safety_num_hidden_layers=1,
+        unfreeze_mm_vision_tower=True,
+        delay_load_vision_tower=False,
+        safety_head_hidden_scale=4.0,
+        pooling_method="mean",  # mean, max, or cls
+        attention_dropout=0.0,  # Add missing attribute for compatibility
+        **kwargs
+    ):
+        # Ensure attention_dropout is in kwargs if not provided
+        if 'attention_dropout' not in kwargs:
+            kwargs['attention_dropout'] = attention_dropout
+
+        super().__init__(**kwargs)
+
+        # Default safety categories if not provided (from original SafeLLaVA)
+        self.safety_categories = safety_categories or [
+            "safe",
+            "gender",
+            "race",
+            "religion",
+            "harassment",
+            "disability_discrimination",
+            "drug_crime",
+            "property_crime",
+            "facial_data",
+            "identity_data",
+            "physical_self_injury",
+            "suicide",
+            "animal_abuse",
+            "obscene_gestures",
+            "physical_altercation",
+            "terrorism",
+            "weapon_related_violence",
+            "sexual_content",
+            "financial_advice",
+            "medical_advice"
+        ]
+
+        self.safety_num_hidden_layers = safety_num_hidden_layers
+        self.unfreeze_mm_vision_tower = unfreeze_mm_vision_tower
+        self.delay_load_vision_tower = delay_load_vision_tower
+        self.safety_head_hidden_scale = safety_head_hidden_scale
+        self.pooling_method = pooling_method
+
+        # self.use_img_safety_meta_token = False
+        self.use_txt_safety_meta_token = False
+        self.use_total_safety_meta_token = False
+
+
+class SafeLlavaLlamaForCausalLM(LlavaLlamaForCausalLM):
+    """
+    SafeLLaVA: A simplified version Uses pooled visual features for safety classification.
+    """
+
+    config_class = SafetyConfig
+    _keys_to_ignore_on_load_unexpected = []  # Don't ignore img_safety_head weights
+
+    def __init__(self, config: SafetyConfig):
+        super().__init__(config)
+
+        # Safety head for image classification (using pooled visual tokens)
+        self.img_safety_head = SafetyMLP(
+            input_size=config.hidden_size,
+            hidden_size=int(config.hidden_size * config.safety_head_hidden_scale),
+            output_size=len(config.safety_categories),
+            safety_num_hidden_layers=config.safety_num_hidden_layers
+        )
+        logging.info("Created img_safety_head for SafeLLaVA")
+
+        # Store pooling method
+        self.pooling_method = config.pooling_method
+
+        # Safety warning template
+        self.safety_warning_template = (
+            "I apologize, but I cannot provide a response as the content appears to be {category}. "
+            "I aim to maintain ethical and safe interactions. "
+            "Please feel free to ask about other topics that do not involve potentially harmful or inappropriate content."
+        )
+
+    def get_model(self):
+        return self.model
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        """
+        Custom from_pretrained to properly load img_safety_head weights.
+        """
+        import os
+        import torch
+        from pathlib import Path
+
+        # Load the model normally first
+        model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # List of original LLaVA model names
+        original_llava_models = [
+            "liuhaotian/llava-v1.5-7b",
+            "liuhaotian/llava-v1.5-13b",
+        ]
+
+        is_original_llava = any(str(pretrained_model_name_or_path).startswith(name) for name in original_llava_models)
+
+        # Load safety head weights for SafeLLaVA models
+        if not is_original_llava:
+            logging.info(f"Detected SafeLLaVA model: {pretrained_model_name_or_path}")
+            model_path = Path(pretrained_model_name_or_path)
+
+            # Handle both local paths and HuggingFace Hub
+            if not model_path.exists():
+                # Try HuggingFace cache
+                from huggingface_hub import snapshot_download
+                try:
+                    model_path = Path(snapshot_download(repo_id=str(pretrained_model_name_or_path)))
+                    logging.info(f"Downloaded from HuggingFace Hub to: {model_path}")
+                except Exception as e:
+                    logging.warning(f"Could not download from Hub: {e}")
+                    return model
+
+            if model_path.exists():
+                # Load safety head weights from safetensors
+                safetensors_index_path = model_path / "model.safetensors.index.json"
+                if safetensors_index_path.exists():
+                    logging.info("Loading safety head weights from safetensors...")
+                    from safetensors.torch import load_file
+                    import json
+
+                    # Load the index file
+                    with open(safetensors_index_path, 'r') as f:
+                        index_data = json.load(f)
+
+                    # Load all safetensors files and collect safety head weights
+                    safety_weights = {}
+                    for weight_map in set(index_data.get('weight_map', {}).values()):
+                        safetensors_file = model_path / weight_map
+                        if safetensors_file.exists():
+                            file_weights = load_file(str(safetensors_file))
+                            # Extract only img_safety_head weights
+                            for key, value in file_weights.items():
+                                if key.startswith('img_safety_head.'):
+                                    safety_weights[key] = value
+
+                    if safety_weights:
+                        logging.info(f"Found {len(safety_weights)} img_safety_head weights")
+                        # Load the weights
+                        missing_keys, unexpected_keys = model.load_state_dict(safety_weights, strict=False)
+                        logging.info("✅ Safety head weights loaded successfully")
+                    else:
+                        logging.warning("⚠️  No img_safety_head weights found in checkpoint")
+                else:
+                    logging.warning(f"No safetensors index found at {safetensors_index_path}")
+            else:
+                logging.warning(f"Model path does not exist: {model_path}")
+
+        return model
+
+    def get_safety_warning(self, unsafe_categories):
+        if len(unsafe_categories) == 1:
+            category_str = f"related to {unsafe_categories[0]}"
+        else:
+            category_str = "related to " + ", ".join(unsafe_categories[:-1]) + f" and {unsafe_categories[-1]}"
+        return self.safety_warning_template.format(category=category_str)
+
+    def pool_visual_tokens(self, hidden_states, input_ids, images):
+        """
+        Pool visual tokens from hidden states.
+
+        Args:
+            hidden_states: Last layer hidden states [batch_size, seq_len, hidden_size]
+            input_ids: Original input token IDs to locate image positions
+            images: Input images tensor
+
+        Returns:
+            Pooled visual features [batch_size, hidden_size]
+        """
+        batch_size = hidden_states.shape[0]
+        device = hidden_states.device
+
+        # If no images, return zeros
+        if images is None:
+            return torch.zeros(batch_size, hidden_states.shape[-1], device=device)
+
+        # Get the number of visual patches
+        vision_tower = self.get_vision_tower()
+        if vision_tower is not None and hasattr(vision_tower, 'config'):
+            # Calculate based on vision config
+            image_size = vision_tower.config.image_size
+            patch_size = vision_tower.config.patch_size
+            num_patches = (image_size // patch_size) ** 2
+        else:
+            num_patches = 576  # Default for CLIP ViT-L/14-336px
+
+        pooled_features = []
+
+        for batch_idx in range(batch_size):
+            try:
+                # Find where IMAGE_TOKEN_INDEX was in the original input
+                if input_ids is not None and batch_idx < input_ids.shape[0]:
+                    image_positions = torch.where(input_ids[batch_idx] == IMAGE_TOKEN_INDEX)[0]
+
+                    if len(image_positions) > 0:
+                        # Visual tokens replace the IMAGE_TOKEN_INDEX
+                        # The actual visual tokens start at this position
+                        start_pos = image_positions[0].item()
+                        end_pos = min(start_pos + num_patches, hidden_states.shape[1])
+
+                        if end_pos > start_pos and (end_pos - start_pos) > 0:
+                            visual_embeddings = hidden_states[batch_idx, start_pos:end_pos]
+
+                            # Apply pooling
+                            if visual_embeddings.shape[0] > 0:
+                                if self.pooling_method == "mean":
+                                    pooled = visual_embeddings.mean(dim=0)
+                                elif self.pooling_method == "max":
+                                    pooled = visual_embeddings.max(dim=0)[0]
+                                elif self.pooling_method == "cls":
+                                    # Use the first visual token
+                                    pooled = visual_embeddings[0]
+                                else:
+                                    pooled = visual_embeddings.mean(dim=0)  # Default to mean
+
+                                pooled_features.append(pooled)
+                            else:
+                                # Empty visual embeddings
+                                pooled_features.append(torch.zeros(hidden_states.shape[-1], device=device))
+                        else:
+                            # Invalid range
+                            pooled_features.append(torch.zeros(hidden_states.shape[-1], device=device))
+                    else:
+                        # No image token found, might be text-only sample
+                        pooled_features.append(torch.zeros(hidden_states.shape[-1], device=device))
+                else:
+                    # No input_ids available
+                    pooled_features.append(torch.zeros(hidden_states.shape[-1], device=device))
+
+            except Exception as e:
+                logging.warning(f"Error pooling visual tokens for batch {batch_idx}: {str(e)}")
+                # Return zero vector on error
+                pooled_features.append(torch.zeros(hidden_states.shape[-1], device=device))
+
+        # Stack all pooled features
+        pooled_features = torch.stack(pooled_features, dim=0)
+        return pooled_features
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        images=None,
+        image_sizes=None,
+        return_dict=None,
+        do_safety=False,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast, SafetyCausalLMOutputWithPast]:
+        """
+        Forward method for SafeLLaVA.
+        When do_safety=True, extracts and pools visual tokens for safety classification.
+        """
+
+        # Store original input_ids for finding image token positions
+        original_input_ids = input_ids.clone() if input_ids is not None else None
+
+        # If do_safety is True, force output_hidden_states to True
+        if do_safety and (output_hidden_states is not True):
+            output_hidden_states = True
+            return_dict = True
+
+        # Prepare inputs for multimodal (handles image embedding)
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images,
+                image_sizes
+            )
+
+        # Call parent's forward method
+        outputs = super(LlavaLlamaForCausalLM, self).forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs
+        )
+
+        # If do_safety=False, just return the outputs
+        if not do_safety:
+            if return_dict is False:
+                return (outputs.loss, outputs.logits, outputs.past_key_values,
+                       outputs.hidden_states, outputs.attentions)
+            return outputs
+
+        # Safety classification using pooled visual tokens
+        hidden_states = outputs.hidden_states[-1]  # Last layer hidden states
+
+        # Check if we have images to process
+        if images is None:
+            # No images, return outputs without safety
+            return outputs
+
+        # Pool visual tokens
+        pooled_visual_features = self.pool_visual_tokens(hidden_states, original_input_ids, images)
+
+        # Pass through safety head
+        img_safety_logits = self.img_safety_head(pooled_visual_features)
+        img_safety_probs = torch.softmax(img_safety_logits, dim=-1)
+
+        # Return results with safety outputs
+        if not return_dict:
+            return (outputs.loss, outputs.logits, outputs.past_key_values,
+                   outputs.hidden_states, outputs.attentions,
+                   img_safety_logits, img_safety_probs)
+
+        return SafetyCausalLMOutputWithPast(
+            loss=outputs.loss,
+            logits=outputs.logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            img_safety_logits=img_safety_logits,
+            img_safety_probs=img_safety_probs,
+            txt_safety_logits=None,  # Not used in Pool version
+            txt_safety_probs=None,
+            total_safety_logits=None,
+            total_safety_probs=None
+        )
+
+
+# Register the model
+AutoConfig.register("safe_llava_llama", SafetyConfig)
+AutoModelForCausalLM.register(SafetyConfig, SafeLlavaLlamaForCausalLM)

safellava/model/llava_arch.py

@@ -0,0 +1,368 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+from .multimodal_encoder.builder import build_vision_tower
+from .multimodal_projector.builder import build_vision_projector
+
+from safellava.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
+from safellava.mm_utils import get_anyres_image_grid_shape
+
+
+class LlavaMetaModel:
+
+    def __init__(self, config):
+        super(LlavaMetaModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            delay_load = getattr(config, 'delay_load_vision_tower', False)
+            self.vision_tower = build_vision_tower(config, delay_load=delay_load)
+            self.mm_projector = build_vision_projector(config)
+
+            if 'unpad' in getattr(config, 'mm_patch_merge_type', ''):
+                self.image_newline = nn.Parameter(
+                    torch.empty(config.hidden_size, dtype=self.dtype)
+                )
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, 'vision_tower', None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def initialize_vision_modules(self, model_args, fsdp=None):
+        vision_tower = model_args.vision_tower
+        mm_vision_select_layer = model_args.mm_vision_select_layer
+        mm_vision_select_feature = model_args.mm_vision_select_feature
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+        mm_patch_merge_type = model_args.mm_patch_merge_type
+
+        self.config.mm_vision_tower = vision_tower
+
+        if self.get_vision_tower() is None:
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [vision_tower]
+            else:
+                self.vision_tower = vision_tower
+        else:
+            if fsdp is not None and len(fsdp) > 0:
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            vision_tower.load_model()
+
+        self.config.use_mm_proj = True
+        self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear')
+        self.config.mm_hidden_size = vision_tower.hidden_size
+        self.config.mm_vision_select_layer = mm_vision_select_layer
+        self.config.mm_vision_select_feature = mm_vision_select_feature
+        self.config.mm_patch_merge_type = mm_patch_merge_type
+
+        if getattr(self, 'mm_projector', None) is None:
+            self.mm_projector = build_vision_projector(self.config)
+
+            if 'unpad' in mm_patch_merge_type:
+                embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
+                self.image_newline = nn.Parameter(
+                    torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std
+                )
+        else:
+            for p in self.mm_projector.parameters():
+                p.requires_grad = True
+
+        if pretrain_mm_mlp_adapter is not None:
+            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')
+            def get_w(weights, keyword):
+                return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k}
+
+            self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector'))
+
+
+def unpad_image(tensor, original_size):
+    """
+    Unpads a PyTorch tensor of a padded and resized image.
+
+    Args:
+    tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format.
+    original_size (tuple): The original size of PIL image (width, height).
+
+    Returns:
+    torch.Tensor: The unpadded image tensor.
+    """
+    original_width, original_height = original_size
+    current_height, current_width = tensor.shape[1:]
+
+    original_aspect_ratio = original_width / original_height
+    current_aspect_ratio = current_width / current_height
+
+    if original_aspect_ratio > current_aspect_ratio:
+        scale_factor = current_width / original_width
+        new_height = int(original_height * scale_factor)
+        padding = (current_height - new_height) // 2
+        unpadded_tensor = tensor[:, padding:current_height - padding, :]
+    else:
+        scale_factor = current_height / original_height
+        new_width = int(original_width * scale_factor)
+        padding = (current_width - new_width) // 2
+        unpadded_tensor = tensor[:, :, padding:current_width - padding]
+
+    return unpadded_tensor
+
+
+class LlavaMetaForCausalLM(ABC):
+
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        vision_tower = self.get_model().get_vision_tower()
+        image_features = vision_tower(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, position_ids, attention_mask, past_key_values, labels,
+        images, image_sizes=None
+    ):
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            return input_ids, position_ids, attention_mask, past_key_values, None, labels
+
+        if type(images) is list or images.ndim == 5:
+            if type(images) is list:
+                images = [x.unsqueeze(0) if x.ndim == 3 else x for x in images]
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            mm_patch_merge_type = getattr(self.config, 'mm_patch_merge_type', 'flat')
+            image_aspect_ratio = getattr(self.config, 'image_aspect_ratio', 'square')
+            if mm_patch_merge_type == 'flat':
+                image_features = [x.flatten(0, 1) for x in image_features]
+            elif mm_patch_merge_type.startswith('spatial'):
+                new_image_features = []
+                for image_idx, image_feature in enumerate(image_features):
+                    if image_feature.shape[0] > 1:
+                        base_image_feature = image_feature[0]
+                        image_feature = image_feature[1:]
+                        height = width = self.get_vision_tower().num_patches_per_side
+                        assert height * width == base_image_feature.shape[0]
+                        if image_aspect_ratio == 'anyres':
+                            num_patch_width, num_patch_height = get_anyres_image_grid_shape(image_sizes[image_idx], self.config.image_grid_pinpoints, self.get_vision_tower().config.image_size)
+                            image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1)
+                        else:
+                            raise NotImplementedError
+                        if 'unpad' in mm_patch_merge_type:
+                            image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
+                            image_feature = image_feature.flatten(1, 2).flatten(2, 3)
+                            image_feature = unpad_image(image_feature, image_sizes[image_idx])
+                            image_feature = torch.cat((
+                                image_feature,
+                                self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device)
+                            ), dim=-1)
+                            image_feature = image_feature.flatten(1, 2).transpose(0, 1)
+                        else:
+                            image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous()
+                            image_feature = image_feature.flatten(0, 3)
+                        image_feature = torch.cat((base_image_feature, image_feature), dim=0)
+                    else:
+                        image_feature = image_feature[0]
+                        if 'unpad' in mm_patch_merge_type:
+                            image_feature = torch.cat((
+                                image_feature,
+                                self.model.image_newline[None].to(image_feature.device)
+                            ), dim=0)
+                    new_image_features.append(image_feature)
+                image_features = new_image_features
+            else:
+                raise ValueError(f"Unexpected mm_patch_merge_type: {self.config.mm_patch_merge_type}")
+        else:
+            image_features = self.encode_images(images)
+
+        if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
+            raise NotImplementedError
+
+        _labels = labels
+        _position_ids = position_ids
+        _attention_mask = attention_mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        else:
+            attention_mask = attention_mask.bool()
+        if position_ids is None:
+            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
+        if labels is None:
+            labels = torch.full_like(input_ids, IGNORE_INDEX)
+
+        _input_ids = input_ids
+        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]
+        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]
+
+        new_input_embeds = []
+        new_labels = []
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
+            if num_images == 0:
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
+                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
+                new_input_embeds.append(cur_input_embeds)
+                new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+
+            image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
+            cur_input_ids_noim = []
+            cur_labels = labels[batch_idx]
+            cur_labels_noim = []
+            for i in range(len(image_token_indices) - 1):
+                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
+                cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
+            split_sizes = [x.shape[0] for x in cur_labels_noim]
+            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
+            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
+            cur_new_input_embeds = []
+            cur_new_labels = []
+
+            for i in range(num_images + 1):
+                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
+                cur_new_labels.append(cur_labels_noim[i])
+                if i < num_images:
+                    cur_image_features = image_features[cur_image_idx]
+                    cur_image_idx += 1
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
+
+            cur_new_input_embeds = [x.to(self.device) for x in cur_new_input_embeds]
+
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
+            cur_new_labels = torch.cat(cur_new_labels)
+
+            new_input_embeds.append(cur_new_input_embeds)
+            new_labels.append(cur_new_labels)
+
+        tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)
+        if tokenizer_model_max_length is not None:
+            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
+            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
+
+        max_len = max(x.shape[0] for x in new_input_embeds)
+        batch_size = len(new_input_embeds)
+
+        new_input_embeds_padded = []
+        new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
+        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)
+        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
+
+        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
+            cur_len = cur_new_embed.shape[0]
+            if getattr(self.config, 'tokenizer_padding_side', 'right') == "left":
+                new_input_embeds_padded.append(torch.cat((
+                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),
+                    cur_new_embed
+                ), dim=0))
+                if cur_len > 0:
+                    new_labels_padded[i, -cur_len:] = cur_new_labels
+                    attention_mask[i, -cur_len:] = True
+                    position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+            else:
+                new_input_embeds_padded.append(torch.cat((
+                    cur_new_embed,
+                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
+                ), dim=0))
+                if cur_len > 0:
+                    new_labels_padded[i, :cur_len] = cur_new_labels
+                    attention_mask[i, :cur_len] = True
+                    position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+
+        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
+
+        if _labels is None:
+            new_labels = None
+        else:
+            new_labels = new_labels_padded
+
+        if _attention_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
+
+        if _position_ids is None:
+            position_ids = None
+
+        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
+
+    def initialize_vision_tokenizer(self, model_args, tokenizer):
+        if model_args.mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+        if model_args.mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+            if model_args.pretrain_mm_mlp_adapter:
+                mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location='cpu')
+                embed_tokens_weight = mm_projector_weights['model.embed_tokens.weight']
+                assert num_new_tokens == 2
+                if input_embeddings.shape == embed_tokens_weight.shape:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
+                elif embed_tokens_weight.shape[0] == num_new_tokens:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
+                else:
+                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")
+        elif model_args.mm_use_im_patch_token:
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = False
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False

safellava/model/multimodal_encoder/__init__.py

@@ -0,0 +1,15 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+"""
+Multimodal encoder module for SafeLLaVA
+"""
+
+from .builder import build_vision_tower
+from .clip_encoder import CLIPVisionTower
+
+__all__ = ['build_vision_tower', 'CLIPVisionTower']

safellava/model/multimodal_encoder/builder.py

@@ -0,0 +1,40 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+"""
+Vision tower builder for SafeLLaVA
+"""
+
+import os
+from .clip_encoder import CLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, custom_vision_tower_path=None, **kwargs):
+    """
+    Build vision tower from config.
+    SafeLLaVA uses CLIPVisionTower.
+    """
+    vision_tower = getattr(
+        vision_tower_cfg,
+        'mm_vision_tower',
+        getattr(vision_tower_cfg, 'vision_tower', None)
+    )
+
+    is_absolute_path_exists = os.path.exists(vision_tower)
+
+    if (is_absolute_path_exists or
+        vision_tower.startswith("openai") or
+        vision_tower.startswith("laion") or
+        "ShareGPT4V" in vision_tower):
+        return CLIPVisionTower(
+            vision_tower,
+            args=vision_tower_cfg,
+            custom_vision_tower_path=custom_vision_tower_path,
+            **kwargs
+        )
+
+    raise ValueError(f'Unknown vision tower: {vision_tower}')

safellava/model/multimodal_encoder/clip_encoder.py

@@ -0,0 +1,127 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+"""
+CLIP Vision Encoder for SafeLLaVA
+
+This is a minimal version containing only what SafeLLaVA needs.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+import logging
+
+from safellava.utils import setup_simple_logging
+
+setup_simple_logging()
+
+
+class CLIPVisionTower(nn.Module):
+    """
+    Basic CLIP vision encoder wrapper for SafeLLaVA.
+    Uses standard CLIP ViT-L/14-336px without modifications.
+    """
+
+    def __init__(self, vision_tower, args, delay_load=False, custom_vision_tower_path=None, **kwargs):
+        super().__init__()
+
+        self.is_loaded = False
+        self.vision_tower_name = vision_tower
+        self.select_layer = args.mm_vision_select_layer
+        self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
+
+        self.use_img_safety_token = False
+
+        # Store custom path if provided (not used in Pool version but needed for compatibility)
+        self.custom_vision_tower_path = custom_vision_tower_path
+
+        if not delay_load:
+            self.load_model()
+        else:
+            self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
+
+    def load_model(self):
+        """Load the CLIP vision model"""
+        if self.is_loaded:
+            return
+
+        logging.info(f"Loading vision tower from: {self.vision_tower_name}")
+
+        # Load standard CLIP model
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)
+        self.vision_tower.requires_grad_(False)
+
+        self.is_loaded = True
+        logging.info("Initialized CLIPVisionModel (no safety tokens)")
+
+    def feature_select(self, image_forward_outs):
+        """Select features from CLIP output"""
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+
+        if self.select_feature == 'patch':
+            # Remove CLS token, keep only patch tokens
+            image_features = image_features[:, 1:]
+        elif self.select_feature == 'cls_patch':
+            # Keep both CLS and patch tokens
+            image_features = image_features
+        else:
+            raise ValueError(f'Unexpected select feature: {self.select_feature}')
+
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        """Forward pass through vision encoder"""
+        if not self.is_loaded:
+            self.load_model()
+
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(
+                    image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                    output_hidden_states=True
+                )
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(
+                images.to(device=self.device, dtype=self.dtype),
+                output_hidden_states=True
+            )
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2

safellava/model/multimodal_projector/__init__.py

@@ -0,0 +1,10 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+from .builder import build_vision_projector, IdentityMap, SimpleResBlock
+
+__all__ = ['build_vision_projector', 'IdentityMap', 'SimpleResBlock']

safellava/model/multimodal_projector/builder.py

@@ -0,0 +1,59 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+import torch
+import torch.nn as nn
+import re
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": 'identity'}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(
+            nn.Linear(channels, channels),
+            nn.GELU(),
+            nn.Linear(channels, channels)
+        )
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+def build_vision_projector(config, delay_load=False, **kwargs):
+    projector_type = getattr(config, 'mm_projector_type', 'linear')
+
+    if projector_type == 'linear':
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == 'identity':
+        return IdentityMap()
+
+    raise ValueError(f'Unknown projector type: {projector_type}')

safellava/utils.py

@@ -0,0 +1,127 @@
+"""
+Based on LLaVA v1.5: https://github.com/haotian-liu/LLaVA
+Modified for SafeLLaVA
+
+Original LLaVA License: Apache License 2.0
+"""
+
+import datetime
+import logging
+import logging.handlers
+import os
+import sys
+
+import requests
+
+from safellava.constants import LOGDIR
+
+server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+handler = None
+
+
+def build_logger(logger_name, logger_filename):
+    global handler
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    if not logging.getLogger().handlers:
+        logging.basicConfig(level=logging.INFO)
+    logging.getLogger().handlers[0].setFormatter(formatter)
+
+    stdout_logger = logging.getLogger("stdout")
+    stdout_logger.setLevel(logging.INFO)
+    sl = StreamToLogger(stdout_logger, logging.INFO)
+    sys.stdout = sl
+
+    stderr_logger = logging.getLogger("stderr")
+    stderr_logger.setLevel(logging.ERROR)
+    sl = StreamToLogger(stderr_logger, logging.ERROR)
+    sys.stderr = sl
+
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    if handler is None:
+        os.makedirs(LOGDIR, exist_ok=True)
+        filename = os.path.join(LOGDIR, logger_filename)
+        handler = logging.handlers.TimedRotatingFileHandler(
+            filename, when='D', utc=True, encoding='UTF-8')
+        handler.setFormatter(formatter)
+
+        for name, item in logging.root.manager.loggerDict.items():
+            if isinstance(item, logging.Logger):
+                item.addHandler(handler)
+
+    return logger
+
+
+class StreamToLogger(object):
+    """Fake file-like stream object that redirects writes to a logger instance."""
+
+    def __init__(self, logger, log_level=logging.INFO):
+        self.terminal = sys.stdout
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def __getattr__(self, attr):
+        return getattr(self.terminal, attr)
+
+    def write(self, buf):
+        temp_linebuf = self.linebuf + buf
+        self.linebuf = ''
+        for line in temp_linebuf.splitlines(True):
+            if line[-1] == '\n':
+                self.logger.log(self.log_level, line.rstrip())
+            else:
+                self.linebuf += line
+
+    def flush(self):
+        if self.linebuf != '':
+            self.logger.log(self.log_level, self.linebuf.rstrip())
+        self.linebuf = ''
+
+
+def disable_torch_init():
+    """Disable the redundant torch default initialization to accelerate model creation."""
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def violates_moderation(text):
+    """Check whether the text violates OpenAI moderation API."""
+    url = "https://api.openai.com/v1/moderations"
+    headers = {"Content-Type": "application/json",
+               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}
+    text = text.replace("\n", "")
+    data = "{" + '"input": ' + f'"{text}"' + "}"
+    data = data.encode("utf-8")
+    try:
+        ret = requests.post(url, headers=headers, data=data, timeout=5)
+        flagged = ret.json()["results"][0]["flagged"]
+    except requests.exceptions.RequestException as e:
+        flagged = False
+    except KeyError as e:
+        flagged = False
+
+    return flagged
+
+
+def pretty_print_semaphore(semaphore):
+    if semaphore is None:
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"
+
+
+def setup_simple_logging():
+    logging.basicConfig(
+        level=logging.INFO,
+        format='[%(asctime)s] [%(levelname)s] [%(filename)s:%(lineno)d:%(funcName)s] %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S,%03d'
+    )

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723