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