FLUX_MODEL_NAME="flux-schnell"
# FLUX_MODEL_NAME="flux-dev"
# HF_TOKEN="[insert HuggingFace token]"

# Install libraries
sudo apt update -y
sudo apt install python3.10-venv -y

# Go to the home directory of Ubuntu
cd /home/ubuntu

# Git clone and set working directory
git clone https://github.com/black-forest-labs/flux.git
cd flux

# Set up the virtual environment
python3 -m venv .flux_venv
source .flux_venv/bin/activate

# Install the packages
pip install -e '.[all]'

# Check if model name is flux-dev
if [ "$FLUX_MODEL_NAME" = "flux-dev" ]; then
  huggingface-cli login --token "$HF_TOKEN"
fi

# Hotfix to allow set default server name to 0.0.0.0 (localhost)
sed -i 's/demo\.launch(share=args\.share)/demo.launch(share=args.share, server_name="0.0.0.0")/' demo_gr.py

# Run the demo on localhost
python3 demo_gr.py --name $FLUX_MODEL_NAME