updated README

README.md

@@ -20,7 +20,7 @@ tags:                            # Add extra tags which would make the repo sear
 
 
 <div align="center">
-  <img src="assets/teaser_camera_ready.svg" width="50%">
+  <img src="assets/teaser_camera_ready.png" width="50%">
 </div>
 
 ## 🚀 Getting Started 
@@ -43,7 +43,7 @@ pip install -e .
 
 Inference of this model requires ~38 GB VRAM on the GPU. Note that the inference runs Blender in EEVEE mode, which runs faster on workstation GPUs like the NVIDIA RTX A6000, compared to data center GPUs like the NVIDIA H100.  
 
-### 🌐 Downloading Pre-Crained checkpoint 
+### 🌐 Download the Pre-Trained Checkpoint 
 
 We use [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) as the base model. To download the SeeThrough3D LoRA checkpoint,
 ```bash
@@ -56,7 +56,7 @@ cd checkpoints
 ### 🤗 Gradio Interface 
 
 It is best to perform inference using the 🤗 Gradio interface, which makes it easy to specify 3D layouts. To launch the interface, run
-```
+```bash
 cd inference 
 conda activate st3d 
 python3 app.py 
@@ -79,15 +79,21 @@ The inference notebook is located at `infer.ipynb`. It is able to load a scene s
 
 ## 🏋 Training  
 
-### 🌐 Downloading the Dataset
+### 🌐 Download the Dataset
+By default, the data is downloaded in the `dataset` directory. To change the download location, edit the `LOCAL_DIR` variable in `dataset/download.py`.  
 
 ```bash
 cd dataset 
-
+conda activate st3d 
+./setup_data.sh 
 ```
 
+We are working on making the data compatible with 🤗 datasets library for ease of visualization and streaming, see [`va1bhavagrawa1/seethrough3d-data`](https://huggingface.co/datasets/va1bhavagrawa1/seethrough3d-data/tree/main) 
+
 ### 🏃 Run Training
 
+Edit `train/train.sh` to specify the downloaded dataset path. 
+
 We train the model for a single epoch at resolution 512, effective batch size of 2 (~25K steps). This requires 2x 80 GB GPUs (one image per GPU). 
 ```bash
 cd train 

dataset/download.py

@@ -0,0 +1,17 @@
+from huggingface_hub import hf_hub_download
+import os
+
+LOCAL_DIR = "."
+
+repo_id = "va1bhavagrawa1/seethrough3d-data"
+file_in_repo = "seethrough3d_data.tar"  # change this
+
+local_path = hf_hub_download(
+    repo_id=repo_id,
+    filename=file_in_repo,
+    repo_type="dataset",
+    local_dir=LOCAL_DIR,                  # current directory
+    local_dir_use_symlinks=False    # makes a real copy (not symlink)
+)
+
+print("Downloaded to:", os.path.abspath(local_path))

dataset/setup_data.sh

@@ -0,0 +1,2 @@
+python3 download.py 
+tar -xvf seethrough3d_data.tar 

inference/config.py

@@ -16,10 +16,9 @@ SAVED_SCENES_DIR = os.path.join(BASE_DIR, "saved_scenes")
 PRETRAINED_MODEL_NAME_OR_PATH = "black-forest-labs/FLUX.1-dev"
 
 # The root directory where LoRA fine-tunes are stored
-LORA_WEIGHTS_ROOT = "/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids"
-
-# The directory containing cached inference embeddings (if used)
-INFERENCE_EMBEDS_DIR = "/archive/vaibhav.agrawal/a-bev-of-the-latents/inference_embeds_flux2"
+LORA_WEIGHTS_ROOT = os.path.join(os.path.dirname(BASE_DIR), "checkpoints") 
+if not os.path.exists(LORA_WEIGHTS_ROOT):
+    raise FileNotFoundError(f"LoRA weights root directory not found: {LORA_WEIGHTS_ROOT}") 
 
 # Available Checkpoint weights (dropdown options)
 CHECKPOINT_NAMES = [

train/train.sh

@@ -1,7 +1,7 @@
 export MODEL_DIR="black-forest-labs/FLUX.1-dev" # your flux path
-export OUTPUT_DIR="/archive/vaibhav.agrawal/a-bev-of-the-latents/checkpoints"  # your save path
+export OUTPUT_DIR="../checkpoints"  # your save path
 export CONFIG="./default_config.yaml"
-export TRAIN_DATA="/archive/vaibhav.agrawal/a-bev-of-the-latents/seethrough3d_data/train.jsonl" # your data jsonl file 
+export TRAIN_DATA="../dataset/seethrough3d_data/train.jsonl" # your data jsonl file 
 export LOG_PATH="$OUTPUT_DIR/log"
 
 accelerate launch --config_file $CONFIG train.py \