diff --git a/README.md b/README.md
index 762f91e..c3f34fc 100644
--- a/README.md
+++ b/README.md
@@ -46,19 +46,20 @@ or make a new separate script like the following
 
 ```python
 import sys
+
 sys.path.append("/path/to/mflux/src")
 
 from flux_1_schnell.config.config import Config
-from flux_1_schnell.models.flux import Flux1Schnell
+from flux_1_schnell.flux import Flux1Schnell
 
 flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
 
 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,
-    )
+   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,
+   )
 )
 
 image.save("image.png")
diff --git a/main.py b/main.py
index bc3c036..7d6385b 100644
--- a/main.py
+++ b/main.py
@@ -4,7 +4,7 @@ import sys
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
 
 from flux_1_schnell.config.config import Config
-from flux_1_schnell.models.flux import Flux1Schnell
+from flux_1_schnell.flux import Flux1Schnell
 
 flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
 
diff --git a/src/flux_1_schnell/models/flux.py b/src/flux_1_schnell/flux.py
similarity index 100%
rename from src/flux_1_schnell/models/flux.py
rename to src/flux_1_schnell/flux.py