Upload Stable Diffusion extraction and PIA rho research code

84468dd verified 3 months ago

7.34 kB

from diffusers import StableDiffusionPipeline
import torch
from typing import Optional, Union, Dict, Any, List, Type
import components
from pytorch_lightning import seed_everything
import logging
from rich.logging import RichHandler
from rich.progress import track
from itertools import chain
import fire
import numpy as np
from dataset import load_member_data
from collections import defaultdict
from torchmetrics.classification import BinaryAUROC, BinaryROC


class EpsGetter(components.EpsGetter):
    def __call__(self, xt: torch.Tensor, condition: torch.Tensor = None, noise_level=None, t: int = None) -> torch.Tensor:
        # t = torch.ones([xt.shape[0]], device=xt.device).long() * t
        return self.model(t, condition, latents=xt)


class MyStableDiffusionPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def prepare_latent(self, img):
        latents = self.vae.encode(img).latent_dist.sample()
        latents = latents * 0.18215
        return latents

    @torch.no_grad()
    def encode_input_prompt(self, prompt, do_classifier_free_guidance=True):
        text_encoder_lora_scale = None
        device = self._execution_device
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            1,
            do_classifier_free_guidance,
            None,
            prompt_embeds=None,
            negative_prompt_embeds=None,
            lora_scale=text_encoder_lora_scale,
        )
        return prompt_embeds

    @torch.no_grad()
    def __call__(
        self,
        t,
        prompt_embeds,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        do_classifier_free_guidance = guidance_scale > 1.0

        # 7. Denoising loop
        # expand the latents if we are doing classifier free guidance
        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents

        # predict the noise residual
        noise_pred = self.unet(
            latent_model_input,
            t,
            encoder_hidden_states=prompt_embeds,
            cross_attention_kwargs=cross_attention_kwargs,
            return_dict=False,
        )[0]

        # perform guidance
        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        return noise_pred

    def get_image(self, latents):
        image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]

        do_denormalize = [True] * image.shape[0]

        image = self.image_processor.postprocess(image, output_type="pil", do_denormalize=do_denormalize)
        return image


def get_FLAGS():

    def FLAGS(x): return x
    FLAGS.T = 1000
    FLAGS.ch = 128
    FLAGS.ch_mult = [1, 2, 2, 2]
    FLAGS.attn = [1]
    FLAGS.num_res_blocks = 2
    FLAGS.dropout = 0.1
    FLAGS.beta_1 = 0.00085
    FLAGS.beta_T = 0.012

    return FLAGS


attackers: Dict[str, Type[components.DDIMAttacker]] = {
    "SecMI": components.SecMIAttacker,
    "PIA": components.PIA,
    "naive": components.NaiveAttacker,
    "PIAN": components.PIAN,
}

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


@torch.no_grad()
def main(attacker_name,
         dataset,
         checkpoint,
         norm_p,
         attack_num=50, interval=10,
         save_logger=None,
         seed=0):
    seed_everything(seed)

    FLAGS = get_FLAGS()

    logger = logging.getLogger()
    logger.disabled = True if save_logger else False
    logger.setLevel(logging.INFO)
    logger.addHandler(RichHandler())

    logger.info("loading model...")
    model = MyStableDiffusionPipeline.from_pretrained(checkpoint, torch_dtype=torch.float32)
    model = model.to(DEVICE)
    # model.eval()

    def attacker_wrapper(attack):
        def wrapper(x, condition=None):
            x = model.prepare_latent(x)
            if 'none' in dataset:
                condition = ['none'] * len(condition)
            if condition is not None:
                condition = model.encode_input_prompt(condition)
            return attack(x, condition)

        return wrapper

    logger.info("loading dataset...")
    _, _, train_loader, test_loader = load_member_data(dataset_name=dataset, batch_size=4)

    attacker = attackers[attacker_name](
        torch.from_numpy(np.linspace(FLAGS.beta_1, FLAGS.beta_T, FLAGS.T)).to(DEVICE), interval, attack_num, EpsGetter(model), lambda x: x * 2 - 1)
    attacker = attacker_wrapper(attacker)

    logger.info("attack start...")
    members, nonmembers = defaultdict(list), defaultdict(list)

    for member, nonmember in track(zip(train_loader, chain(*([test_loader]))), total=len(test_loader)):
        member_condition, nonmenmer_condition = member[1], nonmember[1]
        member, nonmember = member[0].to(DEVICE), nonmember[0].to(DEVICE)

        intermediate_reverse_member, intermediate_denoise_member = attacker(member, member_condition)
        intermediate_reverse_nonmember, intermediate_denoise_nonmember = attacker(nonmember, nonmenmer_condition)

        members[norm_p].append(((intermediate_reverse_member - intermediate_denoise_member).abs() ** norm_p).flatten(2).sum(dim=-1))
        nonmembers[norm_p].append(((intermediate_reverse_nonmember - intermediate_denoise_nonmember).abs() ** norm_p).flatten(2).sum(dim=-1))

        members[norm_p] = [torch.cat(members[norm_p], dim=-1)]
        nonmembers[norm_p] = [torch.cat(nonmembers[norm_p], dim=-1)]

    member = members[0]
    nonmember = nonmembers[0]

    auroc = [BinaryAUROC().cuda()(torch.cat([member[i] / max([member[i].max().item(), nonmember[i].max().item()]), nonmember[i] / max([member[i].max().item(), nonmember[i].max().item()])]), torch.cat([torch.zeros(member.shape[1]).long(), torch.ones(nonmember.shape[1]).long()]).cuda()).item() for i in range(member.shape[0])]
    tpr_fpr = [BinaryROC().cuda()(torch.cat([1 - nonmember[i] / max([member[i].max().item(), nonmember[i].max().item()]), 1 - member[i] / max([member[i].max().item(), nonmember[i].max().item()])]), torch.cat([torch.zeros(member.shape[1]).long(), torch.ones(nonmember.shape[1]).long()]).cuda()) for i in range(member.shape[0])]
    tpr_fpr_1 = [i[1][(i[0] < 0.01).sum() - 1].item() for i in tpr_fpr]
    cp_auroc = auroc[:]
    cp_auroc.sort(reverse=True)
    cp_tpr_fpr_1 = tpr_fpr_1[:]
    cp_tpr_fpr_1.sort(reverse=True)
    print('auc', auroc)
    print('tpr @ 1% fpr', cp_tpr_fpr_1)


if __name__ == '__main__':
    fire.Fire(main)

	from diffusers import StableDiffusionPipeline
	import torch
	from typing import Optional, Union, Dict, Any, List, Type
	import components
	from pytorch_lightning import seed_everything
	import logging
	from rich.logging import RichHandler
	from rich.progress import track
	from itertools import chain
	import fire
	import numpy as np
	from dataset import load_member_data
	from collections import defaultdict
	from torchmetrics.classification import BinaryAUROC, BinaryROC


	class EpsGetter(components.EpsGetter):
	def __call__(self, xt: torch.Tensor, condition: torch.Tensor = None, noise_level=None, t: int = None) -> torch.Tensor:
	# t = torch.ones([xt.shape[0]], device=xt.device).long() * t
	return self.model(t, condition, latents=xt)


	class MyStableDiffusionPipeline(StableDiffusionPipeline):
	@torch.no_grad()
	def prepare_latent(self, img):
	latents = self.vae.encode(img).latent_dist.sample()
	latents = latents * 0.18215
	return latents

	@torch.no_grad()
	def encode_input_prompt(self, prompt, do_classifier_free_guidance=True):
	text_encoder_lora_scale = None
	device = self._execution_device
	prompt_embeds = self._encode_prompt(
	prompt,
	device,
	1,
	do_classifier_free_guidance,
	None,
	prompt_embeds=None,
	negative_prompt_embeds=None,
	lora_scale=text_encoder_lora_scale,
	)
	return prompt_embeds

	@torch.no_grad()
	def __call__(
	self,
	t,
	prompt_embeds,
	prompt: Union[str, List[str]] = None,
	height: Optional[int] = None,
	width: Optional[int] = None,
	guidance_scale: float = 7.5,
	negative_prompt: Optional[Union[str, List[str]]] = None,
	latents: Optional[torch.FloatTensor] = None,
	negative_prompt_embeds: Optional[torch.FloatTensor] = None,
	callback_steps: int = 1,
	cross_attention_kwargs: Optional[Dict[str, Any]] = None,
	):
	# 0. Default height and width to unet
	height = height or self.unet.config.sample_size * self.vae_scale_factor
	width = width or self.unet.config.sample_size * self.vae_scale_factor

	# 1. Check inputs. Raise error if not correct
	self.check_inputs(
	prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
	)

	do_classifier_free_guidance = guidance_scale > 1.0

	# 7. Denoising loop
	# expand the latents if we are doing classifier free guidance
	latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents

	# predict the noise residual
	noise_pred = self.unet(
	latent_model_input,
	t,
	encoder_hidden_states=prompt_embeds,
	cross_attention_kwargs=cross_attention_kwargs,
	return_dict=False,
	)[0]

	# perform guidance
	if do_classifier_free_guidance:
	noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
	noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

	return noise_pred

	def get_image(self, latents):
	image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]

	do_denormalize = [True] * image.shape[0]

	image = self.image_processor.postprocess(image, output_type="pil", do_denormalize=do_denormalize)
	return image


	def get_FLAGS():

	def FLAGS(x): return x
	FLAGS.T = 1000
	FLAGS.ch = 128
	FLAGS.ch_mult = [1, 2, 2, 2]
	FLAGS.attn = [1]
	FLAGS.num_res_blocks = 2
	FLAGS.dropout = 0.1
	FLAGS.beta_1 = 0.00085
	FLAGS.beta_T = 0.012

	return FLAGS


	attackers: Dict[str, Type[components.DDIMAttacker]] = {
	"SecMI": components.SecMIAttacker,
	"PIA": components.PIA,
	"naive": components.NaiveAttacker,
	"PIAN": components.PIAN,
	}

	DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


	@torch.no_grad()
	def main(attacker_name,
	dataset,
	checkpoint,
	norm_p,
	attack_num=50, interval=10,
	save_logger=None,
	seed=0):
	seed_everything(seed)

	FLAGS = get_FLAGS()

	logger = logging.getLogger()
	logger.disabled = True if save_logger else False
	logger.setLevel(logging.INFO)
	logger.addHandler(RichHandler())

	logger.info("loading model...")
	model = MyStableDiffusionPipeline.from_pretrained(checkpoint, torch_dtype=torch.float32)
	model = model.to(DEVICE)
	# model.eval()

	def attacker_wrapper(attack):
	def wrapper(x, condition=None):
	x = model.prepare_latent(x)
	if 'none' in dataset:
	condition = ['none'] * len(condition)
	if condition is not None:
	condition = model.encode_input_prompt(condition)
	return attack(x, condition)

	return wrapper

	logger.info("loading dataset...")
	_, _, train_loader, test_loader = load_member_data(dataset_name=dataset, batch_size=4)

	attacker = attackers[attacker_name](
	torch.from_numpy(np.linspace(FLAGS.beta_1, FLAGS.beta_T, FLAGS.T)).to(DEVICE), interval, attack_num, EpsGetter(model), lambda x: x * 2 - 1)
	attacker = attacker_wrapper(attacker)

	logger.info("attack start...")
	members, nonmembers = defaultdict(list), defaultdict(list)

	for member, nonmember in track(zip(train_loader, chain(*([test_loader]))), total=len(test_loader)):
	member_condition, nonmenmer_condition = member[1], nonmember[1]
	member, nonmember = member[0].to(DEVICE), nonmember[0].to(DEVICE)

	intermediate_reverse_member, intermediate_denoise_member = attacker(member, member_condition)
	intermediate_reverse_nonmember, intermediate_denoise_nonmember = attacker(nonmember, nonmenmer_condition)

	members[norm_p].append(((intermediate_reverse_member - intermediate_denoise_member).abs() ** norm_p).flatten(2).sum(dim=-1))
	nonmembers[norm_p].append(((intermediate_reverse_nonmember - intermediate_denoise_nonmember).abs() ** norm_p).flatten(2).sum(dim=-1))

	members[norm_p] = [torch.cat(members[norm_p], dim=-1)]
	nonmembers[norm_p] = [torch.cat(nonmembers[norm_p], dim=-1)]

	member = members[0]
	nonmember = nonmembers[0]

	auroc = [BinaryAUROC().cuda()(torch.cat([member[i] / max([member[i].max().item(), nonmember[i].max().item()]), nonmember[i] / max([member[i].max().item(), nonmember[i].max().item()])]), torch.cat([torch.zeros(member.shape[1]).long(), torch.ones(nonmember.shape[1]).long()]).cuda()).item() for i in range(member.shape[0])]
	tpr_fpr = [BinaryROC().cuda()(torch.cat([1 - nonmember[i] / max([member[i].max().item(), nonmember[i].max().item()]), 1 - member[i] / max([member[i].max().item(), nonmember[i].max().item()])]), torch.cat([torch.zeros(member.shape[1]).long(), torch.ones(nonmember.shape[1]).long()]).cuda()) for i in range(member.shape[0])]
	tpr_fpr_1 = [i[1][(i[0] < 0.01).sum() - 1].item() for i in tpr_fpr]
	cp_auroc = auroc[:]
	cp_auroc.sort(reverse=True)
	cp_tpr_fpr_1 = tpr_fpr_1[:]
	cp_tpr_fpr_1.sort(reverse=True)
	print('auc', auroc)
	print('tpr @ 1% fpr', cp_tpr_fpr_1)


	if __name__ == '__main__':
	fire.Fire(main)