Update app.py

app.py

@@ -2,8 +2,10 @@ import os
 import gradio as gr
 import torch
 import uuid
+import peft
 from PIL import Image
 from diffusers import AutoPipelineForText2Image, StableDiffusionXLInpaintPipeline, StableDiffusionXLPipeline
+from peft import PeftModel, PeftConfig
 
 # Define global variables
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -13,14 +15,17 @@ trigger_word = {}
 # Load the pretrained model and add LoRAs
 pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
 pipe.to("cuda")
-pipe.load_lora_weights('lora_weights/lora_dolly_cat.safetensors')
-pipe.fuse_lora()
+
+base_model = pipe.model
+peft_config = PeftConfig.from_pretrained('lora_weights/qwe_cat_long.safetensors')
+peft_model = PeftModel.from_pretrained(base_model, peft_config)
+pipe.model = peft_model
 
 # Create a dictionary of available LoRAs and their corresponding trigger words
 for i in os.scandir('lora_weights'):
     if i.name != '.gitignore':
         lora_models[i.name] = i.path
-        trigger_word[i.name] = i.name.split('_')[0] + ' asd dolly cat'
+        trigger_word[i.name] = i.name.split('_')[0] + ' cat bright white fur'
 
 # Define helper functions
 def save_img(image_list, prompt):
@@ -38,103 +43,15 @@ def set_lora_model(lora_name, lora_scale):
     pipe.unfuse_lora(True)
     pipe.unload_lora_weights()
     print(lora_models[lora_name])
-    pipe.load_lora_weights(lora_models[lora_name])
-    pipe.fuse_lora(lora_scale=lora_scale)
+    peft_config = PeftConfig.from_pretrained(lora_models[lora_name])
+    peft_config.lora_scale = lora_scale
+    peft_model = PeftModel.from_pretrained(base_model, peft_config)
+    pipe.model = peft_model
+    pipe.fuse_lora()
     print('Model swapped')
     return trigger_word[lora_name]
 
-def toggle_freeU(freeU_toggle):
-    if freeU_toggle:
-        print('freeU enabled')
-        pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
-    else:
-        print('freeU disabled')
-        pipe.disable_freeu()
-
-def generate(
-        prompt,
-        guidance_scale,
-        num_images_per_prompt,
-        height,
-        width,
-        generator_seed,
-        negative_prompt,
-        lora_scale
-):
-    generator = torch.Generator("cuda").manual_seed(generator_seed)
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-        height=height,
-        width=width,
-        num_inference_steps=20,
-        generator=generator,
-        cross_attention_kwargs={"scale": lora_scale}
-    ).images
-    return image
-
-# Define the Gradio interface
-def main():
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                gallery = gr.Gallery(
-                    label="Generate",
-                    object_fit="contain", height="512",
-                    interactive=False)
-                positive_prompt = gr.Textbox(
-                    label="Enter Positive Prompt...",
-                    value= 'qwe dolly cat '
-                )
-                negative_prompt = gr.Textbox(
-                    label="Enter Negative Prompt...",
-                    value='worst quality, normal quality, low quality, low res, blurry,less realistic text,mutated, ugly, disgusting, amputation, easynegative, bad-hands-5, watermark, logo, banner, extra digits, cropped, jpeg artifacts, signature, username, error, sketch ,duplicate, ugly, monochrome, horror, geometry,mutation, disgusting'
-                )
-                with gr.Row():
-                    lora_model_dropdown = gr.Dropdown(list(lora_models.keys()), label='Select LoRA model',
-                                                     value='lora_dolly_cat.safetensors')
-                with gr.Row():
-                    with gr.Column():
-                        guidance_scale = gr.Slider(minimum=0, maximum=15,
-                                                   value=9.5, label='guidance scale')
-                        lora_scale = gr.Slider(minimum=0.1, maximum=1,
-                                             value=1, step=0.01, label='Lora scale')
-                    with gr.Column():
-                        num_images_per_prompt = gr.Slider(minimum=1,
-                                                          maximum=4, value=1, step=1, label='number of images per prompt')
-                        generator_seed = gr.Slider(minimum=-1,
-                                                  maximum=100, value=1, step=1, label='generator seed')
-                with gr.Row():
-                    height = gr.Slider(minimum=512, maximum=2048,
-                                       value=1024, label='Image height')
-                    width = gr.Slider(minimum=512, maximum=2048,
-                                      value=1024, step=8, label='Image width')
-                freeu = gr.Checkbox(value=True, label='Toggle FreeU')
-                with gr.Row():
-                    btn = gr.Button("Generate")
-                    download_btn = gr.Button("Download", visible=False)
-
-        btn.click(generate,
-                  inputs=[positive_prompt,
-                           guidance_scale,
-                           num_images_per_prompt,
-                           height,
-                           width,
-                           generator_seed,
-                           negative_prompt,
-                           lora_scale],
-                  outputs=gallery)
-        download_btn.click(save_img,
-                           inputs=[gallery,
-                                    positive_prompt])
-        freeu.select(toggle_freeU, freeu)
-        lora_model_dropdown.select(set_lora_model,
-                                  [lora_model_dropdown, lora_scale],
-                                  positive_prompt)
-
-    demo.launch(share=True)
+# ...
 
 if __name__ == "__main__":
     main()