57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
# from https://discuss.huggingface.co/t/generating-and-saving-multiple-images-using-img2img-pipeline/30929
|
|
from diffusers import StableDiffusionImg2ImgPipeline, EulerDiscreteScheduler
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import torch
|
|
import re
|
|
import requests
|
|
|
|
|
|
def slugify(text):
|
|
# remove non-word characters and foreign characters
|
|
text = re.sub(r"[^\w\s]", "", text)
|
|
text = re.sub(r"\s+", "-", text)
|
|
return text
|
|
|
|
model_id = "stabilityai/stable-diffusion-2"
|
|
|
|
images_url = ["https://s3.amazonaws.com/moonup/production/uploads/1675140495576-noauth.png",
|
|
"https://s3.amazonaws.com/moonup/production/uploads/1675032939263-noauth.png",
|
|
"https://s3.amazonaws.com/moonup/production/uploads/1673856328001-noauth.png"]
|
|
|
|
init_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB").resize((758,768)) for url in images_url]
|
|
|
|
prompts = ["beautiful colorful flowr",
|
|
"green city future mountain 3d sunrise skycrapers",
|
|
"rainbow beach, palm trees, neon, miami"]
|
|
|
|
negative_prompts = ["blurry, dark photo, blue",
|
|
"blurry, dark photo, blue",
|
|
"blurry, dark photo, blue"]
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
# Use the Euler scheduler here instead
|
|
scheduler = EulerDiscreteScheduler.from_pretrained(
|
|
model_id, subfolder="scheduler")
|
|
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
|
|
model_id, scheduler=scheduler, torch_dtype=torch.float16)
|
|
pipe = pipe.to(device)
|
|
|
|
DIR_NAME="./images/"
|
|
dirpath = Path(DIR_NAME)
|
|
# create parent dir if doesn't exist
|
|
dirpath.mkdir(parents=True, exist_ok=True)
|
|
|
|
steps = 20
|
|
scale = 9
|
|
num_images_per_prompt = 1
|
|
seed = torch.randint(0, 1000000, (1,)).item()
|
|
generator = torch.Generator(device=device).manual_seed(seed)
|
|
output = pipe(prompts, negative_prompt=negative_prompts, image=init_images, num_inference_steps=steps,
|
|
guidance_scale=scale, num_images_per_prompt=num_images_per_prompt, generator=generator)
|
|
|
|
for idx, (image,prompt) in enumerate(zip(output.images, prompts*num_images_per_prompt)):
|
|
image_name = f'{slugify(prompt)}-{idx}.png'
|
|
image_path = dirpath / image_name
|
|
image.save(image_path)
|