Views
No views yet
@inproceedings{li2024safegen,
author = {Li, Xinfeng and Yang, Yuchen and Deng, Jiangyi and Yan, Chen and Chen, Yanjiao and Ji, Xiaoyu and Xu, Wenyuan},
title = {{SafeGen: Mitigating Sexually Explicit Content Generation in Text-to-Image Models}},
booktitle = {Proceedings of the 2024 {ACM} {SIGSAC} Conference on Computer and Communications Security (CCS)},
year = {2024},
}1#!/bin/bash
2safety_config="MAX"
3prompts_path="<path-to-prompts>"
4model_name="SafeGen_SLD_max"
5image_nums=1
6evaluation_folder="<path-to-save-images>/"${safety_config}_${image_nums}
7model_version="<path-to-SafeGen-Pretrained-Weights>"
8
9python3 SafeGen_SLD_inference.py \
10 --model_name ${model_name} \
11 --model_version ${model_version} \
12 --prompts_path ${prompts_path} \
13 --save_path ${evaluation_folder} \
14 --safety_config ${safety_config} \
15 --num_samples ${image_nums} \
16 --from_case 01'''
2@filename: SafeGen_SLD_inference.py
3@author: Xinfeng Li
4@function: SafeGen can be integrated seamlessly with text-dependent defenses, such as Safe Latent Diffusion (Schramowski et al., CVPR 2023).
5'''
6from diffusers import StableDiffusionPipelineSafe
7from diffusers.pipelines.stable_diffusion_safe import SafetyConfig
8import argparse
9import pandas as pd
10import os
11import torch
12from PIL import Image
13
14device="cuda"
15
16def image_grid(imgs, rows=2, cols=3):
17 w, h = imgs[0].size
18 grid = Image.new("RGB", size=(cols * w, rows * h))
19
20 for i, img in enumerate(imgs):
21 grid.paste(img, box=(i % cols * w, i // cols * h))
22 return grid
23
24row, col = 2, 3
25
26def generate_images(model_name, prompts_path, save_path, device='cuda:0', safety_config="MAX", guidance_scale = 7.5, from_case=0, num_samples=10, model_version="AIML-TUDA/stable-diffusion-safe"):
27 '''
28 Function to generate images from diffusers code
29
30 The program requires the prompts to be in a csv format with headers
31 1. 'case_number' (used for file naming of image)
32 2. 'prompt' (the prompt used to generate image)
33 3. 'seed' (the inital seed to generate gaussion noise for diffusion input)
34
35 Parameters
36 ----------
37 model_name : str
38 name of the model to load.
39 prompts_path : str
40 path for the csv file with prompts and corresponding seeds.
41 save_path : str
42 save directory for images.
43 device : str, optional
44 device to be used to load the model. The default is 'cuda:0'.
45 num_samples : int, optional
46 number of samples generated per prompt. The default is 10.
47 from_case : int, optional
48 The starting offset in csv to generate images. The default is 0.
49
50 Returns
51 -------
52 None.
53 '''
54 pipeline = StableDiffusionPipelineSafe.from_pretrained(model_version)
55 print(pipeline.safety_concept)
56 pipeline = pipeline.to(device)
57
58 df = pd.read_csv(prompts_path)
59
60 folder_path = f'{save_path}/{model_name}'
61 os.makedirs(folder_path, exist_ok=True)
62
63 for _, row in df.iterrows():
64 prompt = [str(row.prompt)]*num_samples
65 case_number = row.case_number
66 if case_number<from_case:
67 continue
68
69 generator = torch.Generator("cuda").manual_seed(int(row.random_seed))
70 if safety_config == "MAX":
71 out_images = pipeline(prompt=prompt, generator=generator, **SafetyConfig.MAX).images
72 elif safety_config == "WEAK":
73 out_images = pipeline(prompt=prompt, generator=generator, **SafetyConfig.WEAK).images
74 elif safety_config == "STRONG":
75 out_images = pipeline(prompt=prompt, generator=generator, **SafetyConfig.STRONG).images
76 elif safety_config == "MEDIUM":
77 out_images = pipeline(prompt=prompt, generator=generator, **SafetyConfig.MEDIUM).images
78
79 for num, im in enumerate(out_images):
80 im.save(f"{folder_path}/{case_number}_{num}.png")
81
82if __name__=='__main__':
83 parser = argparse.ArgumentParser(
84 prog = 'generateImages',
85 description = 'Generate Images using Diffusers Code')
86 parser.add_argument('--model_name', help='name of model', type=str, required=True)
87 parser.add_argument('--model_version', help='path of model', type=str, required=False)
88 parser.add_argument('--prompts_path', help='path to csv file with prompts', type=str, required=True)
89 parser.add_argument('--save_path', help='folder where to save images', type=str, required=True)
90 parser.add_argument('--device', help='cuda device to run on', type=str, required=False, default='cuda:0')
91 parser.add_argument('--guidance_scale', help='guidance to run eval', type=float, required=False, default=7.5)
92 parser.add_argument('--image_size', help='image size used to train', type=int, required=False, default=512)
93 parser.add_argument('--from_case', help='continue generating from case_number', type=int, required=False, default=0)
94 parser.add_argument('--num_samples', help='number of samples per prompt', type=int, required=False, default=1)
95 parser.add_argument('--ddim_steps', help='ddim steps of inference used to train', type=int, required=False, default=100)
96 parser.add_argument('--safety_config', help='safety level [WEAK, MEDIUM, STRONG, MAX]', type=str, required=True, default="MAX")
97 args = parser.parse_args()
98
99 model_name = args.model_name
100 prompts_path = args.prompts_path
101 save_path = args.save_path
102 device = args.device
103 num_samples= args.num_samples
104 from_case = args.from_case
105 safety_config = args.safety_config
106 model_version = args.model_version
107
108 generate_images(model_name, prompts_path, save_path, device=device, safety_config=safety_config,
109 num_samples=num_samples,from_case=from_case, model_version=model_version)<Your_conda_env>/lib/python3.8/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py1# 8. Post-processing
2 image = self.decode_latents(latents)
3
4# 9. Run safety checker
5 # image, has_nsfw_concept, flagged_images = self.run_safety_checker(
6 # image, device, prompt_embeds.dtype, enable_safety_guidance
7 # )
8has_nsfw_concept = None; flagged_images = None
9