Spaces:
Running
Running
File size: 5,019 Bytes
7ccf651 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import gradio as gr
import os
import requests
import json
import base64
from io import BytesIO
from huggingface_hub import login
from PIL import Image
# myip = os.environ["0.0.0.0"]
# myport = os.environ["80"]
myip = "34.219.98.113"
myport=8000
is_spaces = True if "SPACE_ID" in os.environ else False
is_shared_ui = False
from css_html_js import custom_css
from about import (
CITATION_BUTTON_LABEL,
CITATION_BUTTON_TEXT,
EVALUATION_QUEUE_TEXT,
INTRODUCTION_TEXT,
LLM_BENCHMARKS_TEXT,
TITLE,
)
def process_image_from_binary(img_stream):
if img_stream is None:
print("no image binary")
return
image_data = base64.b64decode(img_stream)
image_bytes = BytesIO(image_data)
img = Image.open(image_bytes)
return img
def excute_udiff(diffusion_model_id, concept, steps, attack_id):
print(f"my IP is {myip}, my port is {myport}")
print(f"my input is diffusion_model_id: {diffusion_model_id}, concept: {concept}, steps: {steps}")
response = requests.post('http://{}:{}/udiff'.format(myip, myport),
json={"diffusion_model_id": diffusion_model_id, "concept": concept, "steps": steps, "attack_id": attack_id},
timeout=(10, 1200))
print(f"result: {response}")
# result = result.text[1:-1]
prompt1 = ""
prompt2 = ""
img1 = None
img2 = None
if response.status_code == 200:
response_json = response.json()
print(response_json)
prompt1 = response_json['input_prompt']
prompt2 = response_json['output_prompt']
img1 = process_image_from_binary(response_json['no_attack_img'])
img2 = process_image_from_binary(response_json['attack_img'])
else:
print(f"Request failed with status code {response.status_code}")
return prompt1, prompt2, img1, img2
css = '''
.instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
.arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
#component-4, #component-3, #component-10{min-height: 0}
.duplicate-button img{margin: 0}
#img_1, #img_2, #img_3, #img_4{height:15rem}
#mdStyle{font-size: 0.7rem}
#titleCenter {text-align:center}
'''
with gr.Blocks(css=custom_css) as demo:
gr.HTML(TITLE)
gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
# gr.Markdown("# Demo of UnlearnDiffAtk.")
# gr.Markdown("### UnlearnDiffAtk is an effective and efficient adversarial prompt generation approach for unlearned diffusion models(DMs).")
# # gr.Markdown("####For more details, please visit the [project](https://www.optml-group.com/posts/mu_attack),
# # check the [code](https://github.com/OPTML-Group/Diffusion-MU-Attack), and read the [paper](https://arxiv.org/abs/2310.11868).")
# gr.Markdown("### Please notice that the process may take a long time, but the results will be saved. You can try it later if it waits for too long.")
with gr.Row() as udiff:
with gr.Row():
drop = gr.Dropdown(["Object-Church", "Object-Parachute", "Object-Garbage_Truck","Style-VanGogh",
"Nudity"],
label="Unlearning undesirable concepts")
with gr.Column():
# gr.Markdown("Please upload your model id.")
drop_model = gr.Dropdown(["ESD", "FMN", "SPM"],
label="Unlearned DMs")
# diffusion_model_T = gr.Textbox(label='diffusion_model_id')
# concept = gr.Textbox(label='concept')
# attacker = gr.Textbox(label='attacker')
# start_button = gr.Button("Attack!")
with gr.Column():
atk_idx = gr.Textbox(label="attack index")
with gr.Column():
shown_columns_step = gr.Slider(
0, 100, value=40,
step=1, label="Attack Steps", info="Choose between 0 and 100",
interactive=True,)
with gr.Row() as attack:
with gr.Column(min_width=512):
text_input = gr.Textbox(label="Input Prompt")
orig_img = gr.Image(label="Image Generated by Input Prompt",width=512,show_share_button=False,show_download_button=False)
with gr.Column():
start_button = gr.Button("UnlearnDiffAtk!",size='lg')
with gr.Column(min_width=512):
text_ouput = gr.Textbox(label="Prompt Genetated by UnlearnDiffAtk")
result_img = gr.Image(label="Image Gnerated by Prompt of UnlearnDiffAtk",width=512,show_share_button=False,show_download_button=False)
start_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step, atk_idx], outputs=[text_input, text_ouput, orig_img, result_img], api_name="udiff")
demo.queue().launch(server_name='0.0.0.0') |