From 819a0f7b64516f54c29bcd1895161b36d3544f3d Mon Sep 17 00:00:00 2001
From: ExponentialML <59846140+ExponentialML@users.noreply.github.com>
Date: Sat, 26 Aug 2023 00:03:08 -0700
Subject: [PATCH] More stable LoRA unet training

---
 train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/train.py b/train.py
index 8199964..ec43a94 100644
--- a/train.py
+++ b/train.py
@@ -734,6 +734,8 @@ def finetune_unet(batch, train_encoder=False):
                 is_enabled=True,
                 negation=unet_negation
             )
+            if lora_manager.use_unet_lora:
+                unet.conv_in.requires_grad_(True)
 
         # Convert videos to latent space
         pixel_values = batch["pixel_values"]