add Min-SNR loss to Controlnet flax train script #3016
Conversation
|
The documentation is not available anymore as the PR was closed or merged. |
| @@ -806,6 +812,17 @@ def main(): | |||
| validation_rng, train_rngs = jax.random.split(rng) | |||
| train_rngs = jax.random.split(train_rngs, jax.local_device_count()) | |||
|
|
|||
| def compute_snr(timesteps): | |||
There was a problem hiding this comment.
Could you also add a comment on the reference like here?
|
thank you so much, this is great! Can you kindly also remove the line from the doc that this is only for Pytorch: https://github.com/huggingface/diffusers/blob/main/docs/source/en/training/text2image.mdx#training-with-min-snr-weighting |
|
@kashif I will remove that line once we updated the flax script for text2img (I only updated controlnet in this PR) |
| 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' | ||
| ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' | ||
| ), | ||
| default="wandb", |
There was a problem hiding this comment.
Did we mean to change the default here? (Just asking)
There was a problem hiding this comment.
I just didn't want it to default to something that doesn't exist - cause wandb is the only method I implemented 😅 (and I don't think the other flax training scripts implemented any of these logging methods)
* add wandb team and min-snr loss * make style * apply feedbacks
* add wandb team and min-snr loss * make style * apply feedbacks
* add wandb team and min-snr loss * make style * apply feedbacks
* add wandb team and min-snr loss * make style * apply feedbacks
(a very short) wandb experiment here
https://wandb.ai/yiyixu/train_controlnet_flax?workspace=user-yiyixu