update text encoder codes

app.py

@@ -21,27 +21,38 @@ torch.manual_seed(1)
 counter = 0
 
 #Getting Lora fine-tuned weights
-def monkeypatching(alpha, in_prompt, wt): #, prompt, pipe): finetuned_lora_weights
+def monkeypatching(unet_alpha,texten_alpha, in_prompt, wts): #, prompt, pipe): finetuned_lora_weights
     print("****** inside monkeypatching *******")
     print(f"in_prompt is - {str(in_prompt)}")
     global counter
+    #if model == 'Text-encoder':
+    unet_wt = wts[-2]
+    #else:
+    texten_wt = wts[-1]
+    print(f"UNET weight is = {unet_wt}, Text-encoder weight is = {texten_wt}")
     if counter == 0 :
         #if wt == "./lora_playgroundai_wt.pt" : 
-        monkeypatch_lora(pipe.unet, torch.load(wt)) #finetuned_lora_weights
-        tune_lora_scale(pipe.unet, alpha) #1.00)
+        monkeypatch_lora(pipe.unet, torch.load(unet_wt)) #finetuned_lora_weights
+        monkeypatch_lora(pipe.text_encoder, torch.load(texten_wt), target_replace_module=["CLIPAttention"])  #text-encoder #"./lora/lora_kiriko.text_encoder.pt"
+        #tune_lora_scale(pipe.unet, alpha) #1.00)
+        tune_lora_scale(pipe.unet, unet_alpha)
+        tune_lora_scale(pipe.text_encoder, texten_alpha)
         counter +=1
         #else:
             #monkeypatch_lora(pipe.unet, torch.load("./output_example/lora_weight.pt")) #finetuned_lora_weights
             #tune_lora_scale(pipe.unet, alpha) #1.00)
             #counter +=1
     else :
-        tune_lora_scale(pipe.unet, alpha) #1.00)
-    prompt = "style of hclu, " + str(in_prompt)   #"baby lion"
+        tune_lora_scale(pipe.unet, unet_alpha)
+        tune_lora_scale(pipe.text_encoder, texten_alpha)
+        #tune_lora_scale(pipe.unet, alpha) #1.00)
+    prompt = str(in_prompt) #"style of hclu, " + str(in_prompt)   #"baby lion"
     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, images):
+
+def accelerate_train_lora(steps, images, in_prompt):
     print("*********** inside accelerate_train_lora ***********")
     print(f"images are -- {images}")
     # path can be retrieved by file_obj.name and original filename can be retrieved with file_obj.orig_name
@@ -54,21 +65,27 @@ def accelerate_train_lora(steps, images):
       --pretrained_model_name_or_path={MODEL_NAME}  \
       --instance_data_dir={INSTANCE_DIR} \
       --output_dir={OUTPUT_DIR} \
-      --instance_prompt="style of hclu" \
+      --instance_prompt="{in_prompt}" \
+      --train_text_encoder \
       --resolution=512 \
       --train_batch_size=1 \
       --gradient_accumulation_steps=1 \
       --learning_rate=1e-4 \
+      --learning_rate_text=5e-5 \
+      --color_jitter \
       --lr_scheduler="constant" \
       --lr_warmup_steps=0 \
-      --max_train_steps={int(steps)}') #,shell=True) #30000
+      --max_train_steps={int(steps)}') #10000
     print("*********** completing accelerate_train_lora ***********")
-    print(f"inside output_dir -- {os.listdir(OUTPUT_DIR)}")
+    print(f"files in output_dir -- {os.listdir(OUTPUT_DIR)}")
     #lora_trained_weights = "./output_example/lora_weight.pt" 
-    for file in os.listdir(f"{OUTPUT_DIR}"):
+    files = os.listdir(OUTPUT_DIR)
+    file_list = []
+    for file in files: #os.listdir(OUTPUT_DIR):
         if file.endswith(".pt"):
-            print("weight files is -- ",os.path.join(f"{OUTPUT_DIR}", file))
-            return file 
+            print("weight files are -- ",os.path.join(f"{OUTPUT_DIR}", file))
+            file_list.append(os.path.join(f"{OUTPUT_DIR}", file))
+    return file_list #files[1:] 
     #return f"{OUTPUT_DIR}/*.pt"
     
 with gr.Blocks() as demo:
@@ -77,7 +94,7 @@ with gr.Blocks() as demo:
     gr.HTML("<p>You can skip the queue by duplicating this space and upgrading to gpu in settings: <a style='display:inline-block' href='https://huggingface.co/spaces/ysharma/Low-rank-Adaptation?duplicate=true'><img src='https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14' alt='Duplicate Space'></a></p>")
     gr.Markdown("""<b>NEW!!</b> : I have fine-tuned the SD model for 15,000 steps using 100 PlaygroundAI images and LORA. You can load this trained model using the example component. Load the weight and start using the Space with the Inference button. Feel free to toggle the Alpha value.""")
     gr.Markdown(
-        """**Main Features**<br>- Fine-tune Stable diffusion models twice as faster as dreambooth method by Low-rank Adaptation.<br>- Get insanely small end results, easy to share and download.<br>- Easy to use, compatible with diffusers.<br>- Sometimes even better performance than full fine-tuning<br><br>Please refer to the GitHub repo this Space is based on, here - <a href = "https://github.com/cloneofsimo/lora">LORA</a>. You can also refer to this tweet by AK to quote/retweet/like here on <a href="https://twitter.com/_akhaliq/status/1601120767009513472">Twitter</a>.This Gradio Space is an attempt to explore this novel LORA approach to fine-tune Stable diffusion models, using the power and flexibility of Gradio! The higher number of steps results in longer training time and better fine-tuned SD models.<br><br><b>To use this Space well:</b><br>- First, upload your set of images (4-5), then enter the number of fine-tuning steps, and then press the 'Train LORA model' button. This will produce your fine-tuned model weights.<br>- Enter a prompt, set the alpha value using the Slider (nearer to 1 implies overfitting to the uploaded images), and then press the 'Inference' button. This will produce an image by the newly fine-tuned model.<br><b>Bonus:</b>You can download your fine-tuned model weights from the Gradio file component. The smaller size of LORA models (around 3-4 MB files) is the main highlight of this 'Low-rank Adaptation' approach of fine-tuning.""")
+        """**Main Features**<br>- Fine-tune Stable diffusion models twice as faster as dreambooth method by Low-rank Adaptation.<br>- Get insanely small end results, easy to share and download.<br>- Easy to use, compatible with diffusers.<br>- Sometimes even better performance than full fine-tuning<br><br>Please refer to the GitHub repo this Space is based on, here - <a href = "https://github.com/cloneofsimo/lora">LORA</a>. You can also refer to this tweet by AK to quote/retweet/like here on <a href="https://twitter.com/_akhaliq/status/1601120767009513472">Twitter</a>.This Gradio Space is an attempt to explore this novel LORA approach to fine-tune Stable diffusion models, using the power and flexibility of Gradio! The higher number of steps results in longer training time and better fine-tuned SD models.<br><br><b>To use this Space well:</b><br>- First, upload your set of images (suggested number of images is between 4-9), enter the prompt, enter the number of fine-tuning steps (suggested value is between 2000-4000), and then press the 'Train LORA model' button. This will produce your fine-tuned model weights.<br>- Modify the previous prompt by adding to it (suffix), set the alpha value using the Slider (nearer to 1 implies overfitting to the uploaded images), and then press the 'Inference' button. This will produce an image by the newly fine-tuned UNET and Text-Encoder LORA models.<br><b>Bonus:</b>You can download your fine-tuned model weights from the Gradio file component. The smaller size of LORA models (around 3-4 MB files) is the main highlight of this 'Low-rank Adaptation' approach of fine-tuning.""")
     
     with gr.Row():
         in_images =  gr.File(label="Upload images to fine-tune for LORA", file_count="multiple")
@@ -91,26 +108,28 @@ with gr.Blocks() as demo:
         with gr.Column():
             with gr.Accordion("Advance settings for Training and Inference", open=False):
                 gr.Markdown("Advance settings for a number of Training Steps and Alpha. Set alpha to 1.0 to fully add LORA. If the LORA seems to have too much effect (i.e., overfitting), set alpha to a lower value. If the LORA seems to have too little effect, set the alpha higher. You can tune these two values to your needs.")
-                in_steps = gr.Number(label="Enter the number of training steps", value = 4000)
-                in_alpha = gr.Slider(0.1,1.0, step=0.01, label="Set Alpha level", value=0.5)
-            out_file = gr.File(label="Lora trained model weights" )
+                in_steps = gr.Number(label="Enter the number of training steps", value = 2000)
+                in_alpha_unet = gr.Slider(0.1,1.0, step=0.01, label="Set UNET Alpha level", value=0.5)
+                in_alpha_texten = gr.Slider(0.1,1.0, step=0.01, label="Set Text-Encoder Alpha level", value=0.5)
+                #in_model = gr.Radio(["Text-encoder", "Unet"], label="Select the fine-tuned model for inference", value="Text-encoder", type="value")
+            out_file = gr.File(label="Lora trained model weights", file_count='multiple' )
 
-    gr.Examples(
-        examples=[[0.65, "lion", "./lora_playgroundai_wt.pt" ]],
-        inputs=[in_alpha, in_prompt, out_file],
-        outputs=out_image,
-        fn=monkeypatching,
-        cache_examples=True,)
-    gr.Examples(
-        examples=[[4000, ['./simba1.jpg', './simba2.jpg', './simba3.jpg', './simba4.jpg']]],
-        inputs=[in_steps, in_images],
-        outputs=out_file,
-        fn=accelerate_train_lora,
-        cache_examples=False,
-        run_on_click=False)    
-        
-    b1.click(fn = accelerate_train_lora, inputs=[in_steps, in_images] , outputs=out_file) 
-    b2.click(fn = monkeypatching, inputs=[in_alpha, in_prompt, out_file], outputs=out_image)
+    #gr.Examples(
+    #    examples=[[0.65, 0.6, "lion", ["./lora_playgroundai_wt.pt","./lora_playgroundai_wt.pt"], ],],
+    #    inputs=[in_alpha_unet, in_alpha_texten, in_prompt, out_file ],
+    #    outputs=out_image,
+    #    fn=monkeypatching,
+    #    cache_examples=True,)
+    #gr.Examples(
+    #    examples=[[2500, ['./simba1.jpg', './simba2.jpg', './simba3.jpg', './simba4.jpg'], "baby lion in disney style"]],
+    #    inputs=[in_steps, in_images, in_prompt],
+    #    outputs=out_file,
+    #    fn=accelerate_train_lora,
+    #    cache_examples=False,
+    #    run_on_click=False)    
+    
+    b1.click(fn = accelerate_train_lora, inputs=[in_steps, in_images, in_prompt] , outputs=out_file) 
+    b2.click(fn = monkeypatching, inputs=[in_alpha_unet, in_alpha_texten, in_prompt, out_file,], outputs=out_image)
 
 demo.queue(concurrency_count=3) 
 demo.launch(debug=True, show_error=True)