# Image and Text Embedding Recovery Pipeline

This is a complete workflow for GradCFG.

## Requirements

- Python 3.11
- PyTorch 2.5.1

## Installation

Install dependencies using requirements file:
```bash
pip install -r requirements.txt
```

## Data Preparation

### Dreambooth Dataset

We provide partial Dreambooth dataset with selected categories for quick experimentation. The complete Dreambooth dataset can be downloaded from the official source:

**Partial Dataset (Provided)**
- Location: `dataset`
- Contains selected categories for demonstration
- Ready to use with our pipeline

**Complete Dataset Download**
- Official URL: https://github.com/google/dreambooth
- Full dataset includes 30+ object categories
- Download and extract to `dataset` directory

### Dataset Structure
```
your_dataset_directory/
  ├── dog/           # Category name matches prompt
  │   ├── image1.jpg
  │   ├── image2.jpg
  │   └── ...
  ├── cat/
  └── ...
```

## Workflow

### 1. Download the Model

You can download the model segmind/tiny-sd directly from Hugging Face Hub to local directory `sd_model/tiny-sd`.

### 2. Model Fine-tuning
```bash
python training.py \
  --model_path "sd_model/tiny-sd" \
  --data_dir "dataset" \
  --prompt "dog" \
  --output_dir "ckpts/model_ckpt" \
  --check_epochs 100 \
  --batch_size 1
```
Saves gradients at specified epoch to `ckpts/model_ckpt/{prompt}`

### 3. Reconstruction with GradCFG
```bash
python GradCFG.py \
  --prompt "dog" \
  --model_path "sd_model/tiny-sd" \
  --save_dir "ckpts/GradCFG" \
  --check_epoch 100 \
  --batch_size 1 \
  --num_epochs 4001 \
  --ind_lambda 0.01
```
Outputs reconstructed images and text embeddings in `ckpts/GradCFG/{prompt}`

### 4. Inversion Sampling
```bash
python Inv_Sam.py \
  --prompt "dog" \
  --model_name "sd_model/tiny-sd" \
  --latents_path "ckpts/GradCFG/dog/100_1_100_0.01/dummy_images_epoch4000.pth" \
  --text_path "ckpts/GradCFG/dog/100_1_100_0.01/text_emb_epoch4000.pth" \
  --original_unet_path "ckpts/model_ckpt/dog/epoch_0.pt" \
  --finetuned_unet_path "ckpts/model_ckpt/dog/epoch_200.pt" \
  --output_dir "Inv_Sam_results"
```
Generates final image collage at `Inv_Sam_results/{prompt}.jpg`

### 5. Text Embedding Validation
```bash
python text_generate.py \
  --prompt "dog" \
  --model_dir "sd_model/tiny-sd" \
  --text_embedding_path "ckpts/GradCFG/dog/100_1_100_0.01/text_emb_epoch4000.pth" \
  --output_dir "generated_images" \
  --num_samples 1
```
Validates semantic recovery in `generated_images`

## Directory Structure
```
sd_model/           # Downloaded models
dataset/            # Dataset (partial provided, full download available)
ckpts/
  model_ckpt/       # Training checkpoints
  GradCFG/          # Reconstructed outputs
Inv_Sam_results/    # Final images
generated_images/   # Validation images
```

## Notes

1. GPU requirements:
   - Recommended: 20GB+ VRAM for optimal performance
2. Path placeholders (`{prompt}`) are auto-replaced
3. Adjust hyperparameters as needed per stage
4. Dataset structure: `data_dir/{prompt}/image_files.jpg`
5. Use `--help` flag with any script to see full options
6. For quick start: Use our provided partial Dreambooth dataset in `data/dreambooth/`
7. For full experiments: Download complete Dreambooth dataset from the official source