diff --git a/pan_os_upgrade/__init__.py b/pan_os_upgrade/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/models/__init__.py b/pan_os_upgrade/models/__init__.py
similarity index 100%
rename from models/__init__.py
rename to pan_os_upgrade/models/__init__.py
diff --git a/models/arp_table.py b/pan_os_upgrade/models/arp_table.py
similarity index 100%
rename from models/arp_table.py
rename to pan_os_upgrade/models/arp_table.py
diff --git a/models/assurance_report.py b/pan_os_upgrade/models/assurance_report.py
similarity index 100%
rename from models/assurance_report.py
rename to pan_os_upgrade/models/assurance_report.py
diff --git a/models/content_version.py b/pan_os_upgrade/models/content_version.py
similarity index 100%
rename from models/content_version.py
rename to pan_os_upgrade/models/content_version.py
diff --git a/models/ip_sec_tunnel.py b/pan_os_upgrade/models/ip_sec_tunnel.py
similarity index 100%
rename from models/ip_sec_tunnel.py
rename to pan_os_upgrade/models/ip_sec_tunnel.py
diff --git a/models/license.py b/pan_os_upgrade/models/license.py
similarity index 100%
rename from models/license.py
rename to pan_os_upgrade/models/license.py
diff --git a/models/nics.py b/pan_os_upgrade/models/nics.py
similarity index 100%
rename from models/nics.py
rename to pan_os_upgrade/models/nics.py
diff --git a/models/routes.py b/pan_os_upgrade/models/routes.py
similarity index 100%
rename from models/routes.py
rename to pan_os_upgrade/models/routes.py
diff --git a/models/session_stats.py b/pan_os_upgrade/models/session_stats.py
similarity index 100%
rename from models/session_stats.py
rename to pan_os_upgrade/models/session_stats.py
diff --git a/upgrade.py b/pan_os_upgrade/upgrade.py
similarity index 99%
rename from upgrade.py
rename to pan_os_upgrade/upgrade.py
index c6a83d5..2347728 100644
--- a/upgrade.py
+++ b/pan_os_upgrade/upgrade.py
@@ -111,7 +111,7 @@ class AssuranceOptions:
         },
         "free_disk_space": {
             "description": "Check if a there is enough space on the `/opt/panrepo` volume for downloading an PanOS image.",
-            "log_level": "error",
+            "log_level": "warning",
             "exit_on_failure": False,
         },
         "ha": {
@@ -393,7 +393,6 @@ def parse_arguments() -> Args:
     elif not (arguments["pan_username"] and arguments["pan_password"]):
         logging.error(
             f"{get_emoji('error')} Provide either API key --api-key argument or both --username and --password",
-            file=sys.stderr,
         )
         logging.error(f"{get_emoji('stop')} Halting script.")
 
@@ -1516,7 +1515,7 @@ def perform_reboot(firewall: Firewall, ha_details: Optional[dict] = None) -> Non
                 rebooted = True
             else:
                 logging.info(
-                    f"{get_emoji('working')} Waiting for firewall to reboot and synchronize..."
+                    f"{get_emoji('working')} Firewall is responding to requests but hasn't finished its reboot process..."
                 )
                 time.sleep(30)
 
@@ -1559,6 +1558,20 @@ def main() -> None:
         - If the HA peer state is not synchronized.
         - Upon completion of a dry run.
     """
+
+    # Create necessary directories
+    directories = [
+        "logs",
+        "assurance",
+        "assurance/configurations",
+        "assurance/readiness_checks",
+        "assurance/reports",
+        "assurance/snapshots",
+    ]
+    for dir in directories:
+        ensure_directory_exists(os.path.join(dir, "dummy_file"))
+
+    # Configure logging right after directory setup
     args = parse_arguments()
     configure_logging(args["log_level"])