2 files changed, 66 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..2fa80d6
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+*.png
+
diff --git a/run_flux.py b/run_flux.py
new file mode 100644
index 0000000..c095fce
--- /dev/null
+++ b/run_flux.py
@@ -0,0 +1,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()