# LLaVA-Next for MetaMaterial

## Installation

#### 1. **Navigate to the project directory:**
```bash
cd llava_next_metagen
```

#### 2. **Install the packages:**
```bash
conda create -n llava python=3.10 -y
conda activate llava
pip install --upgrade pip  # Enable PEP 660 support.
pip install -e ".[train]"
```

## Dataset
Download the datasets forreconstruction, inverse design, material understanding, or omni-task models. Specify the dataset paths in the corresponding ``.yaml`` files under the ``data/`` directory.

Download the checkpoints to ``checkpoints/llava_checkpoints_final``:
- https://metagen-datasets.s3.us-east-1.amazonaws.com/v3/workspace/benchmark_inference/llava_checkpoints_final/inverse_design_4target_model.bin
- https://metagen-datasets.s3.us-east-1.amazonaws.com/v3/workspace/benchmark_inference/llava_checkpoints_final/multiview_understanding_model.bin
- https://metagen-datasets.s3.us-east-1.amazonaws.com/v3/workspace/benchmark_inference/llava_checkpoints_final/omni_model.bin
- https://metagen-datasets.s3.us-east-1.amazonaws.com/v3/workspace/benchmark_inference/llava_checkpoints_final/reconstruction_4view_model.bin


## Training and Inference

### 1. Training:

Dowload and unzip training and testing data from https://metagen-datasets.s3.us-east-1.amazonaws.com/anonymized-release/training_data.zip and update paths in data/*.yaml to reflect where you unzip it to.

To train the model, run:
```bash
bash scripts/finetune_peft_metagen.sh ${MID_RUN_NAME} ${DATA_PATH} ${IMG_FOLDER}
```

You may optionally specify ``$MID_RUN_NAME``, ``BASE_MODEL`` and ``VISION_MODEL_VERSION`` in the script.


### 2. Save pretrained model to ``.bin`` file.
The checkpoints are saved at ``checkpoints/${MID_RUN_NAME}/checkpoint-${NUM}/``.

To convert a checkpoint to a ``.bin`` file:
```bash
python checkpoints/${MID_RUN_NAME}/checkpoint-${NUM}/zero_to_fp32.py \
    checkpoints/${MID_RUN_NAME}/checkpoint-${NUM}/ \
    checkpoints/${MID_RUN_NAME}/pytorch_model_${NUM}.bin
```


### 3. Evaluation
To run inference with the system prompts given in the test dataset (recommended):
```bash
bash scripts/evaluation_metagen_newsysprompt.sh ${MID_RUN_NAME} ${CHECKPOINT} ${DEVICE} ${DATA_PATH} ${IMG_FOLDER} 
```

To run inference with the default llava system prompt:
```bash
bash scripts/evaluation_metagen_llavasysprompt.sh ${MID_RUN_NAME} ${CHECKPOINT} ${DEVICE} ${DATA_PATH} ${IMG_FOLDER} 
```

##### Predefined Evaluation Examples:
```bash
# 3.1 Evaluate 4-view reconstruciton model
bash scripts/evaluation_metagen_newsysprompt.sh ${MID_RUN_NAME} checkpoints/llava_checkpoints_final/reconstruction_4view_model.bin ${DEVICE} data/llava_train_4_view_reconstruction.yaml ${IMG_FOLDER}

# 3.2 Evaluate 4-target inverse design model
bash scripts/evaluation_metagen_newsysprompt.sh ${MID_RUN_NAME} checkpoints/llava_checkpoints_final/inverse_design_4target_model.bin ${DEVICE} data/llava_train_4_target_inverse_design.yaml ${IMG_FOLDER}

# 3.3 Evaluate multiview material understanding model
bash scripts/evaluation_metagen_newsysprompt.sh ${MID_RUN_NAME} checkpoints/llava_checkpoints_final/multiview_understanding_model.bin ${DEVICE} data/llava_train_v3_multiview_and_code_material_understanding.yaml ${IMG_FOLDER}

# 3.4 Evaluate omni model
bash scripts/evaluation_metagen_newsysprompt.sh ${MID_RUN_NAME} checkpoints/llava_checkpoints_final/omni_model.bin  ${DEVICE} data/llava_train_v3_omni.yaml ${IMG_FOLDER}
```


### 4. Interactive inference 
```bash
bash scripts/interactive_demo.sh ${MID_RUN_NAME} ${CHECKPOINT} ${DEVICE} ${IMG_FOLDER}
```