update

app.py

@@ -1,21 +1,12 @@
 #https://github.com/huggingface/diffusers/tree/main/examples/dreambooth
-#export MODEL_NAME="stabilityai/stable-diffusion-2-1-base"
-#export INSTANCE_DIR="./data_example"
-#export OUTPUT_DIR="./output_example"
-
-#accelerate launch train_lora_dreambooth.py \
-#  --pretrained_model_name_or_path=$MODEL_NAME  \
-#  --instance_data_dir=$INSTANCE_DIR \
-#  --output_dir=$OUTPUT_DIR \
-#  --instance_prompt="style of sks" \
-#  --resolution=512 \
-#  --train_batch_size=1 \
-#  --gradient_accumulation_steps=1 \
-#  --learning_rate=1e-4 \
-#  --lr_scheduler="constant" \
-#  --lr_warmup_steps=0 \
-#  --max_train_steps=30000
+#export 
+MODEL_NAME="stabilityai/stable-diffusion-2-1-base"
+#export 
+INSTANCE_DIR="./data_example"
+#export 
+OUTPUT_DIR="./output_example"
 
+    
 from diffusers import StableDiffusionPipeline
 from lora_diffusion import monkeypatch_lora, tune_lora_scale
 import torch
@@ -25,7 +16,7 @@ import gradio as gr
 import subprocess
 # If your shell script has shebang,
 # you can omit shell=True argument.
-subprocess.run("./run_lora_db.sh", shell=True)
+#subprocess.run("./run_lora_db.sh", shell=True)
 
 #####
 model_id = "stabilityai/stable-diffusion-2-1-base"
@@ -38,22 +29,39 @@ finetuned_lora_weights = "./lora_weight.pt"
 
 #####
 #my fine tuned weights
-def monkeypatching( alpha): #, prompt, pipe): finetuned_lora_weights
+def monkeypatching(alpha): #, prompt, pipe): finetuned_lora_weights
     monkeypatch_lora(pipe.unet, torch.load(finetuned_lora_weights)) #"./lora_weight.pt")) 
     tune_lora_scale(pipe.unet, alpha) #1.00)
     image = pipe(prompt, num_inference_steps=50, guidance_scale=7).images[0]
     image.save("./illust_lora.jpg") #"./contents/illust_lora.jpg")
     return image
- 
+
+def accelerate_train_lora(steps):
+    accelerate launch "./train_lora_dreambooth.py" \
+      --pretrained_model_name_or_path=MODEL_NAME  \
+      --instance_data_dir=INSTANCE_DIR \
+      --output_dir=OUTPUT_DIR \
+      --instance_prompt="style of sks" \
+      --resolution=512 \
+      --train_batch_size=1 \
+      --gradient_accumulation_steps=1 \
+      --learning_rate=1e-4 \
+      --lr_scheduler="constant" \
+      --lr_warmup_steps=0 \
+      --max_train_steps=steps #30000
+    return
+    
 with gr.Blocks() as demo:
     with gr.Row():
-        in_images =  gr.Image(label="Upload images to fine-tune for LORA")
+        in_images =  gr.File(label="Upload images to fine-tune for LORA", file_count="multiple")
         #in_prompt = gr.Textbox(label="Enter a ")
         in_steps = gr.Number(label="Enter number of steps")
         in_alpha = gr.Slider(0.1,1.0, step=0.01, label="Set Alpha level - higher value has more chances to overfit")
-        b1 = gr.Button(value="Create LORA model")
+        b1 = gr.Button(value="Train LORA model")
+        b2 = gr.Button(value="Inference using LORA model")
     with gr.Row():
         out_image = gr.Image(label="Image generated by LORA model")
-    b1.click(fn = monkeypatching, inputs=in_alpha, outputs=out_image)
+    b1.click(fn = accelerate_train_lora, inputs=in_steps)
+    b2.click(fn = monkeypatching, inputs=in_alpha, outputs=out_image)
 
 demo.launch(debug=True, show_error=True)