# Install the dependencies

!python3 -m virtualenv venv
!. venv/bin/activate

# Upgrade pip to improve performance.
!python3 -m pip install --upgrade pip

# Install PyTorch with the appropriate CUDA support (choose based on the GPU available).
!pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
# Alternatively, install the newer version.
!pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# Clone the DreamCraft3D repository.
!git clone https://github.com/RoenTh/DreamCraft3D.git
%cd DreamCraft3D

# Install additional dependencies.
!pip install ninja
!pip install -r requirements.txt

# Download and set up Zero123 model.
%cd load/zero123
!wget -O stable_zero123.ckpt https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt
%cd ..

# Create directory for Omnidata.
!mkdir -p omnidata
%cd omnidata

# Install tools for downloading from Google Drive.
!sudo apt update
!sudo apt install python3-pip -y
!pip3 install gdown

# Download the Omnidata checkpoints.
!gdown 'https://drive.google.com/uc?id=1Jrh-bRnJEjyMCS7f-WsaFlccfPjJPPHI&confirm=t' # omnidata_dpt_depth_v2.ckpt
!gdown 'https://drive.google.com/uc?id=1wNxVO4vVbDEMEpnAi_jwQObf2MFodcBR&confirm=t' # omnidata_dpt_normal_v2.ckpt
%cd ..
%cd ..

# Install remaining dependencies required for DreamCraft3D.
!pip install carvekit
!pip install ninja
!pip install lightning==2.0.0 omegaconf==2.3.0 jaxtyping typeguard diffusers transformers accelerate opencv-python tensorboard matplotlib imageio imageio[ffmpeg] trimesh bitsandbytes sentencepiece safetensors huggingface_hub libigl xatlas networkx pysdf PyMCubes wandb torchmetrics controlnet_aux
!pip install einops kornia taming-transformers-rom1504 git+https://github.com/openai/CLIP.git
!pip install open3d plotly
!pip install --upgrade setuptools
!pip install mediapipe
!pip install setuptools==69.5.1
!pip install git+https://github.com/ashawkey/envlight.git
!pip install git+https://github.com/KAIR-BAIR/[email protected]
!pip install git+https://github.com/NVlabs/nvdiffrast.git
!pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
!sudo apt-get install libgl1-mesa-glx -y
!pip install mmcv

# Change to DreamCraft3D root directory.
%cd /DreamCraft3D

# Install the Hugging Face diffusers library for DreamBooth training.
!git clone https://github.com/huggingface/diffusers
%cd diffusers
!pip install .

# Return to the root directory.
%cd ..

# Preprocess the reference image for training.
!python preprocess_image_metric3D.py ref.png --recenter

# Create input images for DreamBooth fine-tuning.
!python threestudio/scripts/img_to_mv.py --image_path 'ref_rgba.png' --save_path 'diffusers/examples/dreambooth/ref' --prompt 'a photo of thin turbine-blade' --superres

# Run DreamBooth training with LoRA (Low-Rank Adaptation).
!rm -r ref/.ipynb_checkpoints
!accelerate launch train_dreambooth_lora.py \
    --pretrained_model_name_or_path="DeepFloyd/IF-I-XL-v1.0" \
    --instance_data_dir="ref" \
    --output_dir="blade/results" \
    --instance_prompt="a sks turbine-blade" \
    --resolution=64 \
    --train_batch_size=4 \
    --gradient_accumulation_steps=1 \
    --learning_rate=5e-6 \
    --scale_lr \
    --max_train_steps=1200 \
    --checkpointing_steps=600 \
    --pre_compute_text_embeddings \
    --tokenizer_max_length=77 \
    --text_encoder_use_attention_mask

# Train DreamCraft3D in multiple stages.

# Clean cache directories to free up space.
!rm -rf .cache/*
!rm -rf ~/root/cache/*
!sudo rm -rf /var/cache/*
!sudo rm -rf /root/.cache/*
!sudo rm -rf /tmp/*
!rm -r ~/.local/share/Trash
!df

# Login to Hugging Face to access additional resources.
from huggingface_hub import login
hf_token = "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"
login(token=hf_token)

# Change back to the root directory for training.
%cd /DreamCraft3D

# Define prompt and image path for training.
prompt="a turbine-blade"
image_path="ref_rgba.png"

# Stage 1: DreamCraft3D Coarse NeRF training.
!python launch.py --config configs/dreamcraft3d-coarse-nerf.yaml --train system.prompt_processor.prompt="{prompt}" data.image_path="$image_path" system.guidance.lora_weights_path="diffusers/examples/dreambooth/blade/results"
ckpt = f"outputs/dreamcraft3d-coarse-nerf/{prompt}@LAST/ckpts/last.ckpt"

# Stage 2: DreamCraft3D Coarse NeuS training.
!python launch.py --config configs/dreamcraft3d-coarse-neus.yaml --train system.prompt_processor.prompt="{prompt}" data.image_path="{image_path}" system.weights="{ckpt}" system.guidance.lora_weights_path="diffusers/examples/dreambooth/blade/results"
ckpt = f"outputs/dreamcraft3d-coarse-neus/{prompt}@LAST/ckpts/last.ckpt"

# Stage 3: DreamCraft3D Geometry refinement.
!python launch.py --config configs/dreamcraft3d-geometry.yaml --train system.prompt_processor.prompt="{prompt}" data.image_path="{image_path}" system.geometry_convert_from="{ckpt}" system.guidance.lora_weights_path="diffusers/examples/dreambooth/blade/results"
ckpt = f"outputs/dreamcraft3d-geometry/{prompt}@LAST/ckpts/last.ckpt"

# Stage 4: DreamCraft3D Texture refinement.
!python launch.py --config configs/dreamcraft3d-texture.yaml --train system.prompt_processor.prompt="{prompt}" data.image_path="{image_path}" system.geometry_convert_from="{ckpt}"