You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
57 lines
1.6 KiB
57 lines
1.6 KiB
from diffusers import StableDiffusionPipeline
|
|
import torch
|
|
import random
|
|
import datetime
|
|
|
|
def random_seed(length):
|
|
random.seed()
|
|
min = 10**(length-1)
|
|
max = 9*min + (min-1)
|
|
return random.randint(min, max)
|
|
|
|
device_type = "cuda" # Using AMD GPU with ROCm
|
|
|
|
def load_model():
|
|
model_id = "./models/real-amateur-nudes"
|
|
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None)
|
|
|
|
pipe = pipe.to(device_type)
|
|
return pipe
|
|
|
|
def generate_image(pipe, prompt, seed=None):
|
|
generator = torch.Generator(device=device_type)
|
|
|
|
if seed is not None:
|
|
generator.manual_seed(seed)
|
|
|
|
with torch.no_grad():
|
|
image = pipe(prompt=prompt, num_inference_steps=50, guidance_scale=7.5, generator=generator).images[0]
|
|
return image
|
|
|
|
bullshit_keywords = "hyper detail, cinematic lighting, Canon EOS R3, nikon, f/1.4, ISO 200, 1/160s, 8K, RAW, unedited"
|
|
|
|
def generate(pipe, prompt, seed):
|
|
image = generate_image(pipe, prompt, seed)
|
|
dt = datetime.datetime.now().strftime("%y-%m-%d_%H-%M-%S")
|
|
base_file = 'out/%s_%d' % (dt, seed)
|
|
image_file = '%s.jpg' % base_file
|
|
text_file = '%s.txt' % base_file
|
|
|
|
image.save(image_file)
|
|
with open(text_file, "w") as file:
|
|
file.write("%d\n%s" % (seed, prompt))
|
|
|
|
|
|
def main():
|
|
pipe = load_model()
|
|
prompt = "naked woman, ((shoulder-cut black hair)), %s" % bullshit_keywords
|
|
|
|
seed = 673842166
|
|
generate(pipe, prompt, seed)
|
|
|
|
# for count in range(4):
|
|
# seed = random_seed(9)
|
|
# generate(pipe, prompt, seed)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|