from beam import task_queue, Output, Image, env, Volume
import os
import sys
import requests
if env.is_remote():
     import folder_paths
     sys.path.append(os.getcwd())
     import main


# Function to setup environment and load models
def upload_models():
    # Define the model URLs and their respective directories
    model_urls = [
      #put models here
    ]
    
    # Download models if not already downloaded
    for model_name, filename, directory in model_urls:
        if not os.path.exists(f"{directory}/{filename}"):
            print(os.getcwd())
            os.system(f"wget -c https://huggingface.co/{model_name}/resolve/main/{filename} -P {os.path.join(os.getcwd() , directory)}")



# Import functions or execute code from workflow_api.py

# Step 1: Define environment setup and model loading
@endpoint(
    name="ComfyUI",
    cpu=4,
    memory="32Gi",
    gpu="A10G",
    image=Image(
        python_version="python3.10",
        python_packages=requirements.txt,
        #commands=["pip3 install opencv-python"],
    ),
    volumes=[Volume(name="PUTVOLUMEHERE", mount_path="./VOLUMEPATH)],
    #keep_warm_seconds=0,
    #on_start=upload_models,
)
def generate_video(**inputs):
  while(true):
  main()