File size: 1,372 Bytes
fc4eaf7 2a27c35 fc4eaf7 2a27c35 338ad64 2a27c35 82b5728 b632ef5 2a27c35 cedb7e1 2a27c35 |
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 |
import torch
import gradio as gr
from torchvision import transforms
from diffusers import StableDiffusionPipeline
from model import ResNet, ResidualBlock
from attack import Attack
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1-base"
)
pipe = pipe.to(device)
CLASSES = (
"plane",
"car",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck",
)
def load_classifer(model_path):
# load resnet model
model = ResNet(ResidualBlock, [2, 2, 2])
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
return model
classifer = load_classifer("./models/resnet.ckpt")
attack = Attack(pipe, classifer, device)
def classifer_pred(image):
to_pil = transforms.ToPILImage()
input = attack.transform(to_pil(image[0]))
outputs = classifer(input)
_, predicted = torch.max(outputs, 1)
return CLASSES[predicted[0]]
def run_attack(prompt, epsilon):
image, perturbed_image = attack(prompt, epsilon=epsilon)
pred = classifer_pred(perturbed_image)
return image, pred
demo = gr.Interface(
run_attack,
[gr.Text(), gr.Slider(minimum=0.0, maximum=0.3, value=float)],
[gr.Image(), gr.Text()],
title="Stable Diffused Adversarial Attacks",
)
demo.launch()
|