add lora support

app.py

@@ -157,18 +157,36 @@ vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix",torch_dtype=
 vae.to('cuda')
     
 pipe = StableDiffusionXLPipeline.from_pretrained("John6666/nova-anime-xl-il-v120-sdxl",torch_dtype=torch_dtype,vae=vae)
-
+pipe.load_lora_weights('DervlexVenice/spo_sdxl_4k_p_10ep_lora_webui-base-model-sdxl'
+                           , weight_name='SPO_SDXL_4k_p_10ep_LoRA_webui_510261.safetensors'
+                           , adapter_name="spo")
+pipe.load_lora_weights('DervlexVenice/aesthetic_quality_modifiers_masterpiece-style-illustrious'
+                           , weight_name='Aesthetic_Quality_Modifiers_Masterpiece_929497.safetensors'
+                           , adapter_name="aqm")
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 accelerator = accelerate.Accelerator()
 
-def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps):
+def generate_image_with_steps(prompt
+                              , negative_prompt
+                              , seed
+                              , width
+                              , height
+                              , guidance_scale
+                              , num_inference_steps
+                              , need_lora: bool = False):
     """Helper function to generate image with specific number of steps"""
     scheduler = CustomedUniPCMultistepScheduler.from_config(pipe.scheduler.config
                                                             , solver_order = 2 if num_inference_steps==8 else 1
                                                             ,denoise_to_zero = False
                                                             , use_afs=True)
+    if not need_lora:
+        pipe.set_adapters(["spo", "aqm"], adapter_weights=[0.0, 0.0])
+    elif num_inference_steps > 6:
+        pipe.set_adapters(["spo", "aqm"], adapter_weights=[0.7, 0.7])
+    else:
+        pipe.set_adapters(["spo", "aqm"], adapter_weights=[0.25, 0.25])
     pipe.scheduler = scheduler
     pipe.to('cuda')
     with torch.no_grad():
@@ -240,6 +258,7 @@ def infer(
     resolution,
     guidance_scale,
     num_inference_steps,
+    need_lora,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -249,7 +268,7 @@ def infer(
     width, height = map(int, resolution.split('x'))
     
     # Generate image with selected steps
-    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps)
+    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, need_lora=need_lora)
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, final_sigmas_type="sigma_min")
     # Generate image with 50 steps for high quality
     negative_prompts = '(worst quality:2), (low quality:2), (normal quality:2), bad anatomy, bad proportions, poorly drawn face, poorly drawn hands, missing fingers, extra limbs, blurry, pixelated, distorted, lowres, jpeg artifacts, watermark, signature, text, (deformed:1.5), (bad hands:1.3), overexposed, underexposed, censored, mutated, extra fingers, cloned face, bad eyes'
@@ -343,6 +362,11 @@ with gr.Blocks() as demo:
                     label="Number of inference steps",
                 )
 
+            need_lora = gr.Checkbox(
+                label="Use LoRA adapters",
+                value=True,
+            )
+
         gr.Examples(examples=examples, inputs=[prompt])
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -355,6 +379,7 @@ with gr.Blocks() as demo:
             resolution,
             guidance_scale,
             num_inference_steps,
+            need_lora,
         ],
         outputs=[result, result_20_steps, seed],
     )