Delete app.py

app.py

@@ -1,253 +0,0 @@
-import json
-import os
-import os.path as osp
-import threading
-
-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
-
-
-UNCONSTRAINED_ATTACK_IMAGE_PATH = 'unconstrained_attack_images/adversarial_'
-CONSTRAINED_ATTACK_IMAGE_PATH = 'adversarial_qna_images/adv_image_'
-MODEL_PATH = "liuhaotian/llava-v1.5-13b"
-
-TEXT_SAFETY_PATCHES = {
-    "optimized": "text_patch_optimized",
-    "heuristic": "text_patch_heuristic"
-}
-
-IMAGE_SAFETY_PATCHES = {
-    "default": "safety_patch.pt"
-}
-
-
-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 = tokenizer = image_processor = context_len = my_generator = None
-
-def load_model_async(model_path, model_name):
-    global tokenizer, model, image_processor, context_len, loaded_model_name, my_generator
-    print(f"Loading {model_name} model ... ")
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
-    if "llava" in model_name.lower():
-        loaded_model_name = "LLaVA"
-    else:
-        raise NotImplementedError
-
-    my_generator = generator.Generator(model=model, tokenizer=tokenizer)
-
-threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()
-
-print('>>> Initializing Models')
-
-prompts = rtp_read(osp.join('harmful_corpus/rtp_prompts.jsonl'))
-
-# out_unprotected: responses without the safety patch
-out, out_unprotected = [], []
-text_prompt = 'Please extend the following sentence: %s'
-
-
-def generate_answer(image, user_message: str, requested_model_name: str,
-                    image_safety_patch_type: str, text_safety_patch_type: str,
-                    ):
-    global tokenizer, model, image_processor, context_len, loaded_model_name, my_generator
-
-    with open(TEXT_SAFETY_PATCHES[text_safety_patch_type], 'r') as file:
-        text_safety_patch = file.read().rstrip()
-
-    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 ... ")
-
-            threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()
-            my_generator = generator.Generator(model=model, tokenizer=tokenizer)
-
-        # load a randomly-sampled unconstrained attack image as Image object
-        if isinstance(image, str):
-            image = load_image(image)
-
-        # 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()
-
-        user_message_unprotected = user_message
-        if text_safety_patch != None:
-            if text_safety_patch_type == "optimal":
-                # use the below for optimal text safety patch
-                user_message = text_safety_patch + '\n' + user_message
-
-            elif text_safety_patch_type == "heuristic":
-                # use the below for heuristic text safety patch
-                user_message += '\n' + text_safety_patch
-            else:
-                raise ValueError(f"Invalid safety patch type: {user_message}")
-
-        text_prompt_template_unprotected = prompt_wrapper.prepare_text_prompt(text_prompt % user_message_unprotected)
-        prompt_unprotected = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template_unprotected,
-                                                   device=model.device)
-
-        text_prompt_template = prompt_wrapper.prepare_text_prompt(text_prompt % user_message)
-        prompt = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template, device=model.device)
-
-        response_unprotected = my_generator.generate(prompt_unprotected, image).replace("[INST]", "").replace("[/INST]",
-                                                                                                              "").replace(
-            "[SYS]", "").replace("[/SYS/]", "").strip()
-
-        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, "")
-
-        response_unprotected = response_unprotected.replace(text_safety_patch, "")
-
-        print(" -- [Unprotected] continuation: ---")
-        print(response_unprotected)
-        print(" -- [Protected] continuation: ---")
-        print(response)
-
-        out.append({'prompt': user_message, 'continuation': response})
-        out_unprotected.append({'prompt': user_message, 'continuation': response_unprotected})
-
-        return response, response_unprotected
-
-
-def get_list_of_examples():
-    global rtp
-    examples = []
-
-    # Use the first 3 prompts for constrained attack
-    for i, prompt in enumerate(prompts[:3]):
-        image_num = np.random.randint(25)  # Randomly select an image number
-        image_path = f'{CONSTRAINED_ATTACK_IMAGE_PATH}{image_num}.bmp'
-
-        examples.append(
-            [image_path, prompt]
-        )
-
-    # Use the 3-6th prompts for unconstrained attack
-    for i, prompt in enumerate(prompts[3:6]):
-        image_num = np.random.randint(25)  # Randomly select an image number
-        image_path = f'{UNCONSTRAINED_ATTACK_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(image_path: str, user_message: 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")
-    response, response_unprotected = generate_answer(image_path, user_message, model_selector.value, image_patch_selector.value,
-                    text_patch_selector.value)
-
-    return response, response_unprotected
-
-
-with gr.Blocks(css=css) as demo:
-    state = get_empty_state()
-    all_components = []
-
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            """# 🦙LLaVAGuard🔥<br>
-    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=[
-                gr.Textbox(label="With Safety Patches"),
-                gr.Textbox(label="NO Safety Patches")
-            ])
-
-
-# Launch the demo
-demo.launch()