WAT-ai-AA/stable-diffused-adversarial-attacks

1import torch2import gradio as gr3from torchvision import transforms4from diffusers import StableDiffusionPipeline5from model import ResNet, ResidualBlock6from attack import Attack7 8device = "cuda" if torch.cuda.is_available() else "cpu"9 10pipe = StableDiffusionPipeline.from_pretrained(11    "stabilityai/stable-diffusion-2-1-base"12)13pipe = pipe.to(device)14 15CLASSES = (16    "plane",17    "car",18    "bird",19    "cat",20    "deer",21    "dog",22    "frog",23    "horse",24    "ship",25    "truck",26)27 28 29def load_classifer(model_path):30    # load resnet model31    model = ResNet(ResidualBlock, [2, 2, 2])32    model.load_state_dict(torch.load(model_path, map_location=device))33    model.eval()34    return model35 36 37classifer = load_classifer("./models/resnet.ckpt")38attack = Attack(pipe, classifer, device)39 40 41def classifer_pred(image):42    to_pil = transforms.ToPILImage()43    input = attack.transform(to_pil(image[0]))44    outputs = classifer(input)45    _, predicted = torch.max(outputs, 1)46    return CLASSES[predicted[0]]47 48 49def run_attack(prompt, epsilon):50    image, perturbed_image = attack(prompt, epsilon=epsilon)51    pred = classifer_pred(perturbed_image)52    return image, pred53 54 55demo = gr.Interface(56    run_attack,57    [gr.Text(), gr.Slider(minimum=0.0, maximum=0.3, value=float)],58    [gr.Image(), gr.Text()],59    title="Stable Diffused Adversarial Attacks",60)61demo.launch()62