import json import os.path as osp import gradio as gr import numpy as np import torch from llava.mm_utils import get_model_name_from_path from llava.model.builder import load_pretrained_model from llava_utils import prompt_wrapper, generator from utils import normalize, denormalize, load_image # BASE_DIR = "/workingdir/soh337/llavaguard" BASE_DIR = "/Users/ahren/Workspace/Multimodal/llavaguard" IMAGE_PATH = 'unconstrained_attack_images/adversarial_' MODEL_PATH = "/workingdir/models_hf/liuhaotian/llava-v1.5-13b" TEXT_SAFETY_PATCHES = { "optimized": "text_patch_optimized", "heuristic": "text_patch_heuristic" } IMAGE_SAFETY_PATCHES = { "default": "safety_patch.bmp" } def rtp_read(text_file): dataset = [] lines = open(text_file).read().split("\n") for li in lines: obj = json.loads(li) if obj['challenging']: dataset.append(obj['prompt']['text']) return dataset model = loaded_model_name = None print('>>> Initializing Models') prompts = rtp_read(osp.join(BASE_DIR, 'harmful_corpus/rtp_prompts.jsonl')) out = [] def generate_answer(user_message: str, image_path: str, requested_model_name: str, image_safety_patch_type: str, text_safety_patch_type: str, ): global loaded_model_name text_safety_patch = TEXT_SAFETY_PATCHES[text_safety_patch_type] image_safety_patch = IMAGE_SAFETY_PATCHES[image_safety_patch_type] if requested_model_name == "LLaVA": if requested_model_name == loaded_model_name: print(f"{requested_model_name} model already loaded.") else: print(f"Loading {requested_model_name} model ... ") model_name = get_model_name_from_path(MODEL_PATH) tokenizer, model, image_processor, context_len = load_pretrained_model(MODEL_PATH, None, model_name) loaded_model_name = requested_model_name my_generator = generator.Generator(model=model, tokenizer=tokenizer) # load a randomly-sampled unconstrained attack image as Image object image = load_image(image_path) # transform the image using the visual encoder (CLIP) of LLaVA 1.5; the processed image size would be PyTorch tensor whose shape is (336,336). image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda() if image_safety_patch != None: # make the image pixel values between (0,1) image = normalize(image) # load the safety patch tensor whose values are (0,1) safety_patch = torch.load(image_safety_patch).cuda() # apply the safety patch to the input image, clamp it between (0,1) and denormalize it to the original pixel values safe_image = denormalize((image + safety_patch).clamp(0, 1)) # make sure the image value is between (0,1) print(torch.min(image), torch.max(image), torch.min(safe_image), torch.max(safe_image)) else: safe_image = image model.eval() if text_safety_patch != None: # use the below for optimal text safety patch # user_message = text_safety_patch + '\n' + user_message # use the below for heuristic text safety patch user_message += '\n' + text_safety_patch text_prompt_template = prompt_wrapper.prepare_text_prompt(text_prompt % user_message) print(text_prompt_template) prompt = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template, device=model.device) response = my_generator.generate(prompt, safe_image).replace("[INST]", "").replace("[/INST]", "").replace( "[SYS]", "").replace("[/SYS/]", "").strip() if text_safety_patch != None: response = response.replace(text_safety_patch, "") print(" -- continuation: ---") print(response) out.append({'prompt': user_message, 'continuation': response}) def get_list_of_examples(): global rtp examples = [] for i, prompt in enumerate(prompts[:3]): # Use the first 3 prompts for simplicity image_num = np.random.randint(25) # Randomly select an image number image_path = f'{IMAGE_PATH}{image_num}.bmp' examples.append( [image_path, prompt] ) return examples css = """#col-container {max-width: 90%; margin-left: auto; margin-right: auto; display: flex; flex-direction: column;} #header {text-align: center;} #col-chatbox {flex: 1; max-height: min(750px, 100%);} #label {font-size: 2em; padding: 0.5em; margin: 0;} .message {font-size: 1.2em;} .message-wrap {max-height: min(700px, 100vh);} """ def get_empty_state(): # TODO: Not sure what this means return gr.State({"arena": None}) examples = get_list_of_examples() # Define a function to update inputs based on selected example def update_inputs(example_id): selected_example = examples[int(example_id)] return selected_example['image_path'], selected_example['text'] model_selector, image_patch_selector, text_patch_selector = None, None, None def process_text_and_image(user_message: str, image_path: str): global model_selector, image_patch_selector, text_patch_selector print(f"User Message: {user_message}") # print(f"Text Safety Patch: {safety_patch}") print(f"Image Path: {image_path}") print(model_selector.value) # generate_answer(user_message, image_path, "LLaVA", "heuristic", "default") generate_answer(user_message, image_path, model_selector.value, image_patch_selector.value, text_patch_selector.value) with gr.Blocks(css=css) as demo: state = get_empty_state() all_components = [] with gr.Column(elem_id="col-container"): gr.Markdown( """# 🦙LLaVAGuard🔥
Safeguarding your Multimodal LLM **[Project Homepage](#)**""", elem_id="header", ) # example_selector = gr.Dropdown(choices=[f"Example {i}" for i, e in enumerate(examples)], # label="Select an Example") with gr.Row(): model_selector = gr.Dropdown(choices=["LLaVA"], label="Model", info="Select Model", value="LLaVA") image_patch_selector = gr.Dropdown(choices=["default"], label="Image Patch", info="Select Image Safety " "Patch", value="default") text_patch_selector = gr.Dropdown(choices=["heuristic", "optimized"], label="Text Patch", info="Select " "Text " "Safety " "Patch", value="heuristic") image_and_text_uploader = gr.Interface( fn=process_text_and_image, inputs=[gr.Image(type="pil", label="Upload your image", interactive=True), gr.Textbox(placeholder="Input a question", label="Your Question"), ], examples=examples, outputs=['text']) # # Set the action for the generate button # @demo.events(generate_button) # def handle_generation(image, question, model, image_patch, text_patch): # generate_answer(question, image, model, text_patch, image_patch) # Launch the demo demo.launch()