runpod_setup

gistfile1.py

#!/usr/bin/env python3
from pathlib import Path
import subprocess
import os
import sys
import json

STARTUP_SH="start.sh"
REFID="4501dd11ea6060740fd764b2f6006173c68b28ae"
WORKDIR=Path("/workspace")

def install_pip_packages(*packages):
    print("Installing Python packages")
    if Path("requirements.txt").exists():
        print("Found requirements.txt file")
        subprocess.run([
            f"pip install --default-timeout=100 -r requirements.txt"
        ],shell=True)
    pkg_str = " ".join(packages)
    print(f"Installing {pkg_str}")
    subprocess.run([
        f"pip install {pkg_str}"
    ],shell=True)
    print("Purging cache for space")
    subprocess.run([
        f'pip cache purge'
    ],shell=True)

def setup_viton():
    subprocess.run([f"wget https://github.com/akhildevelops/catvton-flux/archive/{REFID}.tar.gz -O catvton-flux.tar.gz"],shell=True)
    subprocess.run([f"tar -xvf catvton-flux.tar.gz"],shell=True)
    os.chdir(WORKDIR/f"catvton-flux-{REFID}")
    install_pip_packages("protobuf","huggingface_hub[cli,hf_transfer]")
    if not os.environ.get("HF_TOKEN"):
        print("HF_TOKEN is not found in the environment, exiting.")
        sys.exit(1)
    subprocess.run([
        f"huggingface-cli download black-forest-labs/FLUX.1-dev"
    ],shell=True)
    subprocess.run([
        f"huggingface-cli download xiaozaa/catvton-flux-alpha"
    ],shell=True)



def main():
    print("These are the env variables")
    print(os.environ)
    WORKDIR.mkdir(parents=True,exist_ok=True)        
    print(f"Changing to Working Directory: {WORKDIR}")
    os.chdir(WORKDIR)
    print("Setting up Viton")
    setup_viton()
    os.chdir("/")
    startup_script_path = Path(STARTUP_SH)
    if startup_script_path.exists():
        print("Found start up script, running")
        subprocess.run([f"./{STARTUP_SH}"],shell=True)
    else:
        print("Didn't found startup script")
    print(f"Successfuly setup environment in {WORKDIR}")
    subprocess.run(["sleep infinity"],shell=True)
    
    

if __name__=="__main__":
    main()