| 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 |
| |
| 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 |
| |
| 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") |
|
|
|
|
| |
