diff --git a/README.md b/README.md index 06bebbeb..967d43a5 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f Run the provided [main.py](main.py) ``` -python main.py +python main.py --prompt "luxury food photograph" --steps 2 ``` or make a new separate script like the following @@ -135,4 +135,3 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots - FLUX Dev implementation - LoRA adapters -- Command line args \ No newline at end of file diff --git a/main.py b/main.py index 5cf39bb6..1da111ec 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,7 @@ import os import sys +import argparse +import time sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) @@ -7,16 +9,34 @@ from flux_1_schnell.flux import Flux1Schnell from flux_1_schnell.post_processing.image_util import ImageUtil -flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell") +def main(): + parser = argparse.ArgumentParser(description='Generate an image based on a prompt.') + parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.') + parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".') + parser.add_argument('--model', type=str, default="black-forest-labs/FLUX.1-schnell", help='The model to use. Default is "black-forest-labs/FLUX.1-schnell".') + parser.add_argument('--seed', type=int, default=0, help='Entropy Seed (Default is time-based random-seed)') + parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)') + parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)') + parser.add_argument('--steps', type=int, default=4, help='Inference Steps') -image = flux.generate_image( - seed=3, - prompt="Luxury food photograph of a birthday cake. In the middle it has three candles shaped like letters spelling the word 'MLX'. It has perfect lighting and a cozy background with big bokeh and shallow depth of field. The mood is a sunset balcony in tuscany. The photo is taken from the side of the cake. The scene is complemented by a warm, inviting light that highlights the textures and colors of the ingredients, giving it an appetizing and elegant look.", - config=Config( - num_inference_steps=2, - height=768, - width=1360, + args = parser.parse_args() + + seed = args.seed or time.time() + + flux = Flux1Schnell(args.model) + + image = flux.generate_image( + seed=args.seed, + prompt=args.prompt, + config=Config( + num_inference_steps=args.steps, + height=args.height, + width=args.width, + ) ) -) -ImageUtil.save_image(image, "image.png") + image.save(args.output) + +if __name__ == '__main__': + main() +