48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
import os
|
|
import argparse
|
|
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from dotenv import load_dotenv
|
|
|
|
from cloudflare_api import text_to_image, refine_text_prompt
|
|
from generate_model_local import image_to_3d_api, image_to_3d_subprocess
|
|
|
|
load_dotenv()
|
|
|
|
PIPELINE_FOLDER = os.environ["PIPELINE_FOLDER"]
|
|
|
|
|
|
def get_timestamp():
|
|
return datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
|
|
|
|
|
|
def main():
|
|
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()
|
|
|
|
input_prompt = args.prompt
|
|
print(f"Input prompt: {input_prompt}")
|
|
|
|
refine_prompt = os.environ["REFINE_PROMPT"] == "1"
|
|
if refine_prompt:
|
|
image_generation_prompt = refine_text_prompt(input_prompt)
|
|
print(f"Refined prompt: {image_generation_prompt}")
|
|
else:
|
|
image_generation_prompt = input_prompt
|
|
|
|
timestamp = get_timestamp()
|
|
pipeline_folder = Path(PIPELINE_FOLDER)
|
|
image_path = pipeline_folder / "images" / f"{timestamp}.jpg"
|
|
text_to_image(image_generation_prompt, image_path)
|
|
print(f"Generated image file: {image_path}")
|
|
model_path = pipeline_folder / "models" / timestamp
|
|
image_to_3d_api(image_path, model_path)
|
|
#model_file_path = model_path / "0" / "mesh.glb"
|
|
print(f"Generated 3D model file: {model_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|