diff --git a/README.md b/README.md
index 967d43a..dca8aef 100644
--- a/README.md
+++ b/README.md
@@ -37,12 +37,29 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
    ```
 ### Generating an image
 
-Run the provided [main.py](main.py) 
+Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so:
+
 ```
-python main.py --prompt "luxury food photograph" --steps 2
+python main.py --prompt "Luxury food photograph" --steps 2 --seed 2
 ```
 
-or make a new separate script like the following
+#### Full list of Command-Line Arguments 
+
+- **`--prompt`** (required, `str`): Text description of the image to generate.
+
+- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename.
+
+- **`--model`** (optional, `str`, default: `"black-forest-labs/FLUX.1-schnell"`): Model to use for generation.
+
+- **`--seed`** (optional, `int`, default: `0`): Seed for random number generation. Default is time-based.
+
+- **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.
+
+- **`--width`** (optional, `int`, default: `1024`): Width of the output image in pixels.
+
+- **`--steps`** (optional, `int`, default: `4`): Number of inference steps.
+
+Or make a new separate script like the following
 
 ```python
 import sys
@@ -60,6 +77,8 @@ image = flux.generate_image(
    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,
    )
 )
 
diff --git a/main.py b/main.py
index 1da111e..97b3f39 100644
--- a/main.py
+++ b/main.py
@@ -9,6 +9,7 @@ from flux_1_schnell.config.config import Config
 from flux_1_schnell.flux import Flux1Schnell
 from flux_1_schnell.post_processing.image_util import ImageUtil
 
+
 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.')
@@ -21,7 +22,7 @@ def main():
 
     args = parser.parse_args()
 
-    seed = args.seed or time.time()
+    seed = args.seed or int(time.time())
 
     flux = Flux1Schnell(args.model)
 
@@ -35,8 +36,8 @@ def main():
         )
     )
 
-    image.save(args.output)
+    ImageUtil.save_image(image, "image.png")
+
 
 if __name__ == '__main__':
     main()
-