Skip to content

Instantly share code, notes, and snippets.

@kg6zjl
Created February 24, 2023 07:18
Show Gist options
  • Select an option

  • Save kg6zjl/1aa19941de765234f9d3d98332cbfa18 to your computer and use it in GitHub Desktop.

Select an option

Save kg6zjl/1aa19941de765234f9d3d98332cbfa18 to your computer and use it in GitHub Desktop.
#!/usr/bin/env python3
from diffusers import StableDiffusionPipeline
import torch
import random
import string
import time
def save_image(image, name, count):
image.save(f"C:/stable-diffusion/images/{name}-{count}.png")
def random_string(length):
letters = string.ascii_lowercase
return ''.join(random.choice(letters) for i in range(length))
def image_generator(prompt, model_id, image_count = 4,):
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)
pipe.enable_attention_slicing()
pipe = pipe.to("cuda")
# create a random string for the filename
name = random_string(10)
for count in range(image_count):
image = pipe(prompt).images[0]
save_image(image, name, count)
prompts = [
"seasons changing while kids play in the street of at an abandoned wartorn village in eastern Europe"
]
image_count = 5
for prompt in prompts:
image_generator(prompt, "runwayml/stable-diffusion-v1-5", image_count)
time.sleep(1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment