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
|
import torch
import os
import base64
import argparse
from diffusers import FluxPipeline
from typing import Tuple
import uuid
def load_flux():
pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
pipeline.enable_model_cpu_offload()
pipeline.vae.enable_slicing()
pipeline.vae.enable_tiling()
return pipeline
def generate_image(pipeline, prompt, height=1024, width=1024, guideance_scale=0, num_images_per_prompt=1, num_inference_steps=50):
images = pipeline(
prompt=prompt,
guidance_scale=guideance_scale,
height=height,
width=width,
max_sequence_length=256,
num_inference_steps=num_inference_steps,
num_images_per_prompt=num_images_per_prompt
).images
return images
def generate_random_string(length=16) -> str:
return str(uuid.uuid4())
def parse_dimensions(dim_str: str) -> Tuple[int, int]:
try:
width, height = map(int, dim_str.split(':'))
return (width, height)
except ValueError:
raise argparse.ArgumentError('Dimensions must be in format width:height')
def main():
try:
args = parser.parse_args()
pipeline = load_flux()
pipeline.to(torch.float16)
width, height = args.size
images = generate_image(pipeline, prompt=args.prompt, width=width, height=height, guideance_scale=args.guideance_scale, num_images_per_prompt=args.number)
for image in images:
filename = generate_random_string()
image.save(f"{filename}.png")
except KeyboardInterrupt:
print('\nExiting early...')
exit(0)
parser = argparse.ArgumentParser(description="Generate some A.I. images", epilog="All done!")
parser.add_argument("-n", "--number", type=int, default=1, help="the number of images you want to generate")
parser.add_argument("-o", "--output", type=str, default="image", help="the name of the output image")
parser.add_argument("-p", "--prompt", type=str, required=True, help="the prompt")
parser.add_argument("-gs", "--guideance-scale", type=float, default=0)
parser.add_argument("--size", type=parse_dimensions, default="1024:1024")
if __name__ == "__main__":
main()
|
