Train model on edge2shoes dataset

[jigyasa27]

Hello
I am trying to use the code to train model on edge2shoes dataset (https://www.kaggle.com/datasets/balraj98/edges2shoes-dataset).
I have used the default setting in the code and training command is -

python src/train_pix2pix_turbo.py -pretrained_model_name_or_path="stabilityai/sd-turbo" -output_dir="pix2pix_turbo_final/edge2shoes" --dataset_folder="/data/edge2shoes" --resolution=512 --train_batch_size=2 --max_train_steps=10000 --learning_rate=5e-6 --enable_xformers_memory_efficient_attention --viz_freq 100 --checkpointing_steps 500 ---eval_freq 500 --report_to "wandb" --tracker_project_name "edgeshoes"

After training for 85000 iterations also, I am not able to get good images
Val FID at 21000 iter = 114 (all image are black in color)
Val FID at 85000 iter = 283 (images appear distorted)

In the text prompt, I have used same prompt for train and val images
"realistic photo of a shoe, detailed and colored image"

• Can you guide how to improve the results on this dataset.

• I was trying to use "edge_to_image" model name while training, however this results in error
  "Pix2Pix_Turbo" object has no attribute 'target_modules'
  Since there was no command line argument for model name, I had modified the function call
  net_pix2pix = Pix2Pix_Turbo("edge_to_image")