passing user prompt to pipeline script as cmd argument

3d-generation-pipeline/start_pipeline.py

@@ -1,6 +1,7 @@
 import os
 import base64
 import requests
+import argparse
 from datetime import datetime
 from dotenv import load_dotenv
 
@@ -66,10 +67,44 @@ def refine_text_prompt(prompt):
     data = response.json()
     return data["result"]["response"]
 
+def image_to_3d(image_path, output_path):
+    import subprocess
+
+    MODEL_FOLDER = r"D:\users\henrisel\stable-fast-3d"
+
+    venv_python = MODEL_FOLDER + r"\.venv\Scripts\python.exe"
+    script_path = MODEL_FOLDER + r"\run.py"
+
+    args = [image_path, "--output-dir", output_path]
+    command = [venv_python, script_path] + args
+
+    try:
+        # Run the subprocess
+        result = subprocess.run(command, capture_output=True, text=True)
+
+        # Print output and errors
+        print("STDOUT:\n", result.stdout)
+        print("STDERR:\n", result.stderr)
+        print("Return Code:", result.returncode)
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
 def main():
-    user_prompt = "Give epic sword"
+    parser = argparse.ArgumentParser(description="Text to 3D model pipeline")
+    parser.add_argument("--prompt", type=str, required=True, help="User text prompt")
+    args = parser.parse_args()
+
+    user_prompt = args.prompt
     print(f"User prompt: {user_prompt}")
     refined_prompt = refine_text_prompt(user_prompt)
     print(f"Refined prompt: {refined_prompt}")
     image_path = text_to_image(refined_prompt, get_timestamp())
     print(f"Generated image file: {image_path}")
+    output_path = r"\models"
+    image_to_3d(image_path, output_path)
+    print(f"Generated 3D model file: {output_path}")
+
+
+if __name__ == "__main__":
+    main()