| import torch |
| import torch.nn.functional as F |
|
|
| prompt_dataset = [ |
| "Portrait of an astronaut in space, detailed starry background, reflective helmet,", |
| "Painting of a floating island with giant clock gears, populated with mythical creatures,", |
| "Landscape of a Japanese garden in autumn, with a bridge over a koi pond,", |
| "Painting representing the sound of jazz music, using vibrant colors and erratic shapes,", |
| "Painting of a modern smartphone with classic art pieces appearing on the screen,", |
| "Battle scene with futuristic robots and a golden palace in the background,", |
| "Scene of a bustling city market with different perspectives of people and stalls,", |
| "Scene of a ship sailing in a stormy sea, with dramatic lighting and powerful waves,", |
| "Portraint of a female botanist surrounded by exotic plants in a greenhouse,", |
| "Painting of an ancient castle at night, with a full moon, gargoyles, and shadows,", |
| ] |
|
|
| style_dataset = [ |
| "Art Nouveau", |
| "Romantic", |
| "Cubist", |
| "Baroque", |
| "Pop Art", |
| "Abstract", |
| "Impressionist", |
| "Surrealist", |
| "Renaissance", |
| "Pointillism", |
| ] |
|
|
|
|
|
|
| class attack_mixin: |
| def __call__( |
| self, |
| latents: torch.Tensor, |
| timesteps: torch.Tensor, |
| encoder_hidden_states: torch.Tensor, |
| unet: torch.nn.Module, |
| target_tensor: torch.Tensor, |
| noise_scheduler |
| ): |
| raise NotImplementedError |
| |
| class AdvDM(attack_mixin): |
| """ |
| This attack aims to maximize the training loss of diffusion model |
| """ |
| def __call__( |
| self, |
| latents: torch.Tensor, |
| noise: torch.Tensor, |
| timesteps: torch.Tensor, |
| encoder_hidden_states: torch.Tensor, |
| unet: torch.nn.Module, |
| text_encoder: torch.nn.Module, |
| input_ids, |
| target_tensor: torch.Tensor, |
| noise_scheduler |
| ): |
| noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) |
|
|
| |
| encoder_hidden_states = text_encoder(input_ids)[0] |
|
|
| |
| model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample |
|
|
| |
| if noise_scheduler.config.prediction_type == "epsilon": |
| target = noise |
| elif noise_scheduler.config.prediction_type == "v_prediction": |
| target = noise_scheduler.get_velocity(latents, noise, timesteps) |
| else: |
| raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") |
|
|
| unet.zero_grad() |
| text_encoder.zero_grad() |
| loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") |
|
|
| |
| if target_tensor is not None: |
| xtm1_pred = torch.cat( |
| [ |
| noise_scheduler.step( |
| model_pred[idx : idx + 1], |
| timesteps[idx : idx + 1], |
| noisy_latents[idx : idx + 1], |
| ).prev_sample |
| for idx in range(len(model_pred)) |
| ] |
| ) |
| xtm1_target = noise_scheduler.add_noise(target_tensor, noise, timesteps - 1) |
| loss = loss - F.mse_loss(xtm1_pred, xtm1_target) |
|
|
| return loss |
| |
| class LatentAttack(attack_mixin): |
| """ |
| This attack aims to minimize the l2 distance between latent and target_tensor |
| """ |
| def __call__( |
| self, |
| latents: torch.Tensor, |
| timesteps: torch.Tensor=None, |
| encoder_hidden_states: torch.Tensor=None, |
| unet: torch.nn.Module=None, |
| target_tensor: torch.Tensor=None, |
| noise_scheduler=None |
| ): |
| if target_tensor == None: |
| raise ValueError("Need a target tensor for pre-attack") |
| loss = - F.mse_loss(latents, target_tensor, reduction="mean") |
| return loss |
