# WaterCycleDiffusion

**Repository Path**: wanghaoupc/WaterCycleDiffusion

## Basic Information

- **Project Name**: WaterCycleDiffusion
- **Description**: WaterCycleDiffusion.
- **Primary Language**: Python
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-06-27
- **Last Updated**: 2025-08-05

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

#### Prerequisites
1. Python 3.10
2. CUDA 11.5

#### Compilation
Install all the python dependencies using pip:
        
        pip install -r requirements.txt

#### Train
        accelerate config
        export NCCL_P2P_DISABLE=1
        accelerate launch --main_process_port 29501 train.py \
        --pretrained_model_name_or_path="stabilityai/sd-turbo" \
        --output_dir="output/VTF/U2N" \
        --dataset_folder "data/U2N" \
        --train_img_prep "resize_286_randomcrop_256x256_hflip" --val_img_prep "no_resize" \
        --learning_rate="1e-5" --max_train_steps=25000 \
        --train_batch_size=1 --gradient_accumulation_steps=1 \
        --report_to "wandb" --tracker_project_name "U2N_v2" \
        --enable_xformers_memory_efficient_attention --validation_steps 250 \
        --lambda_gan 0.5 --lambda_idt 1 --lambda_cycle 1

#### Demo
        python src/demo.py --model_path "output/VTF/U2N/checkpoints/model_6001.pkl" \
            --input_image "data/VTF/test_A/000001.jpg" \
            --prompt "This is a water-free image" --direction "a2b" \
            --output_dir "outputs" --image_prep "no_resize"
    
#### Acknowledgment
Our work uses the Cyclegan as the base model with the following [LICENSE](https://github.com/GaParmar/img2img-turbo/blob/main/LICENSE).