File size: 2,889 Bytes
e225ed6 |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 |
from datasets import load_dataset
from diffusers import StableDiffusionPipeline
import torch
import os
import json
from PIL import Image
import argparse
def parse_args():
parser = argparse.ArgumentParser(description="Generate images from I2P dataset")
parser.add_argument("--output_dir", type=str, help="Output directory")
parser.add_argument("--model_path", type=str, help="Path to model checkpoint", default="CompVis/stable-diffusion-v1-4")
parser.add_argument("--seed", type=int, help="Seed for random number generator", default=0)
parser.add_argument("--prompt", type=str, help="Prompt for image generation")
parser.add_argument("--mode", type=str, help="Mode for image generation", choices=["train","test"], default="train")
parser.add_argument("--num_train_images", type=int, help="Number of images to generate for training", default=1000)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if(args.mode == "test"):
metadata = []
os.makedirs(args.output_dir, exist_ok=True)
pipe = StableDiffusionPipeline.from_pretrained(args.model_path, safety_checker=None, torch_dtype=torch.float16).to(device)
for i in range(args.num_train_images):
while True:
output = pipe(prompt=args.prompt)
image = output.images[0]
nsfw = output.nsfw_content_detected
#check if nsfw is a list
if isinstance(nsfw, list):
nsfw = nsfw[0]
if not nsfw:
break
image.save(os.path.join(args.output_dir, f"{args.prompt}_{i}.png"))
elif(args.mode == "train"):
metadata = []
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.join(args.output_dir, "train"), exist_ok=True)
pipe = StableDiffusionPipeline.from_pretrained(args.model_path, safety_checker=None, torch_dtype=torch.float16).to(device)
for i in range(args.num_train_images):
while True:
output = pipe(prompt=args.prompt)
image = output.images[0]
nsfw = output.nsfw_content_detected
#check if nsfw is a list
if isinstance(nsfw, list):
nsfw = nsfw[0]
if not nsfw:
break
image.save(os.path.join(args.output_dir, "train", f"{args.prompt}_{i}.png"))
metadata.append({"file_name": f"train/{args.prompt}_{i}.png", "text": args.prompt})
with open(os.path.join(args.output_dir, 'metadata.jsonl'), 'w') as f:
for m in metadata:
f.write(json.dumps(m) + "\n")
|