app.py

import gradio as gr
import os
import av
import base64
import cv2
import torch
import time
import numpy as np
import requests
import json
from huggingface_hub import login


myip = os.environ["myip"]
myport = os.environ["myport"]


is_spaces = True if "SPACE_ID" in os.environ else False

is_shared_ui = False


def excute_udiff(model_id, concept, attacker):
    print(f"my IP is {myip}, my port is {myport}")
    result = requests.post('http://{}:{}/udiff'.format(myip, myport), json={"model_id": model_id, "concept": concept, "attacker": attacker})
    
    return f"The attack prompt is: {result}"


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=css) as demo:
    gr.Markdown("## Unlearn Diffusions Attack")
    gr.Markdown("Check if your model is safe.")

    with gr.Row() as udiff:
        with gr.Column():
            # gr.Markdown("Please upload your model id.")
            model_id = gr.Textbox(label='model_id')
            concept = gr.Textbox(label='concept')
            attacker = gr.Textbox(label='attacker')

            start_button = gr.Button("Attack!")

        with gr.Column():
            result = gr.Textbox(label="attack prompt")

    start_button.click(fn=excute_udiff, inputs=[model_id, concept, attacker], outputs=result, api_name="udiff")


# demo.queue(default_enabled=False, api_open=False, max_size=5).launch(debug=True, show_api=False)
demo.queue(concurrency_count=5).launch(debug=True, server_name='0.0.0.0')