DeltaVR/3d-generation-pipeline/start_pipeline.py

63 lines
2.1 KiB
Python

import argparse
import asyncio
import logging
import time
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv
from generate_image_local import text_to_image_invoke_ai
from generate_model_local import image_to_3d_api
load_dotenv()
def get_timestamp():
return datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
def setup_logger(base_folder, timestamp):
log_dir = base_folder / Path("logs")
log_dir.mkdir(parents=True, exist_ok=True)
logging.basicConfig(
filename=log_dir / f"{timestamp}.log",
level=logging.INFO,
force=True
)
async 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}")
image_generation_prompt = input_prompt + ", single object, plain neutral background, high-resolution, sharp details, soft studio lighting, front and side visible, realistic style, true-to-scale"
pipeline_folder = Path(__file__).resolve().parent
timestamp = get_timestamp()
setup_logger(pipeline_folder, timestamp)
time_checkpoint = time.time()
image_path = pipeline_folder / "images" / f"{timestamp}.jpg"
# TODO: use Invoke AI or Cloudflare, depending on env var
#text_to_image_cloudflare(image_generation_prompt, image_path)
await text_to_image_invoke_ai(image_generation_prompt, image_path)
image_generation_time = time.time() - time_checkpoint
time_checkpoint = time.time()
logging.info(f"Image generation time: {round(image_generation_time, 1)} s")
print(f"Generated image file: {image_path}")
model_path = pipeline_folder / "models" / timestamp
model_file = image_to_3d_api(image_path, model_path)
model_generation_time = time.time() - time_checkpoint
logging.info(f"Model generation time: {round(model_generation_time, 1)} s")
print(f"Generated 3D model file: {model_file}")
if __name__ == "__main__":
asyncio.run(main())