Update app.py

app.py

@@ -123,6 +123,8 @@ pipe.load_ip_adapter_instantid(face_adapter)
 pipe.set_ip_adapter_scale(0.8)
 zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
 zoe.to(device)
+
+original_pipe = copy.deepcopy(pipe)
 pipe.to(device)
 
 last_lora = ""
@@ -202,10 +204,58 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 )
                 del weights_sd
                 del lora_model
-@spaces.GPU    
+@spaces.GPU
+def generate_image(prompt, negative, face_emb, face_image, image_strength, images, guidance_scale, face_strength, depth_control_scale, last_lora, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index):
+    if last_lora != repo_name:
+        if(last_fused):
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+        pipe.load_lora_weights(loaded_state_dict)
+        pipe.fuse_lora(lora_scale)
+        last_fused = True
+        is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
+        if(is_pivotal):
+            #Add the textual inversion embeddings from pivotal tuning models
+            text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
+            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
+            state_dict_embedding = load_file(embedding_path)
+            print(state_dict_embedding)
+            try:
+                pipe.unload_textual_inversion()
+                pipe.load_textual_inversion(state_dict_embedding["clip_l"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+                pipe.load_textual_inversion(state_dict_embedding["clip_g"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+            except:
+                pipe.unload_textual_inversion()
+                pipe.load_textual_inversion(state_dict_embedding["text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+                pipe.load_textual_inversion(state_dict_embedding["text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+
+    print("Processing prompt...")
+    conditioning, pooled = compel(prompt)
+    if(negative):
+        negative_conditioning, negative_pooled = compel(negative)
+    else:
+        negative_conditioning, negative_pooled = None, None
+    print("Processing image...")
+    image = pipe(
+        prompt_embeds=conditioning,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=negative_conditioning,
+        negative_pooled_prompt_embeds=negative_pooled,
+        width=1024,
+        height=1024,
+        image_embeds=face_emb,
+        image=face_image,
+        strength=1-image_strength,
+        control_image=images,
+        num_inference_steps=20,
+        guidance_scale = guidance_scale,
+        controlnet_conditioning_scale=[face_strength, depth_control_scale],
+    ).images[0]
+    return image
+    
 def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, progress=gr.Progress(track_tqdm=True)):
     global last_lora, last_merged, last_fused, pipe
-
+    selected_state_index = selected_state.index
     face_image = center_crop_image_as_square(face_image)
     try:
         face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
@@ -216,7 +266,7 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
         raise gr.Error("No face found in your image. Only face images work here. Try again")
         
     for lora_list in lora_defaults:
-        if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
+        if lora_list["model"] == sdxl_loras[selected_state_index]["repo"]:
             prompt_full = lora_list.get("prompt", None)
             if(prompt_full):
                 prompt = prompt_full.replace("<subject>", prompt)
@@ -224,7 +274,7 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     
     print("Prompt:", prompt) 
     if(prompt == ""):
-        prompt = "A person"
+        prompt = "a person"
     #prepare face zoe
     with torch.no_grad():
         image_zoe = zoe(face_image)
@@ -239,15 +289,15 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     #    else:
     #         selected_state.index *= -1
         #sdxl_loras = sdxl_loras_new
-    print("Selected State: ", selected_state.index)
-    print(sdxl_loras[selected_state.index]["repo"])
+    print("Selected State: ", selected_state_index)
+    print(sdxl_loras[selected_state_index]["repo"])
     if negative == "":
         negative = None
 
     if not selected_state:
         raise gr.Error("You must select a LoRA")
-    repo_name = sdxl_loras[selected_state.index]["repo"]
-    weight_name = sdxl_loras[selected_state.index]["weights"]
+    repo_name = sdxl_loras[selected_state_index]["repo"]
+    weight_name = sdxl_loras[selected_state_index]["weights"]
     
     full_path_lora = state_dicts[repo_name]["saved_name"]
     loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
@@ -255,53 +305,8 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     print("Last LoRA: ", last_lora)
     print("Current LoRA: ", repo_name)
     print("Last fused: ", last_fused)
-    if last_lora != repo_name:
-        if(last_fused):
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
-        pipe.load_lora_weights(loaded_state_dict)
-        pipe.fuse_lora(lora_scale)
-        last_fused = True
-        is_pivotal = sdxl_loras[selected_state.index]["is_pivotal"]
-        if(is_pivotal):
-            #Add the textual inversion embeddings from pivotal tuning models
-            text_embedding_name = sdxl_loras[selected_state.index]["text_embedding_weights"]
-            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
-            state_dict_embedding = load_file(embedding_path)
-            print(state_dict_embedding)
-            try:
-                pipe.unload_textual_inversion()
-                pipe.load_textual_inversion(state_dict_embedding["clip_l"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
-                pipe.load_textual_inversion(state_dict_embedding["clip_g"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
-            except:
-                pipe.unload_textual_inversion()
-                pipe.load_textual_inversion(state_dict_embedding["text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
-                pipe.load_textual_inversion(state_dict_embedding["text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
-    
-    print("Processing prompt...")
-    conditioning, pooled = compel(prompt)
-    if(negative):
-        negative_conditioning, negative_pooled = compel(negative)
-    else:
-        negative_conditioning, negative_pooled = None, None
-    print("Processing image...")
-    
-    image = pipe(
-        prompt_embeds=conditioning,
-        pooled_prompt_embeds=pooled,
-        negative_prompt_embeds=negative_conditioning,
-        negative_pooled_prompt_embeds=negative_pooled,
-        width=1024,
-        height=1024,
-        image_embeds=face_emb,
-        image=face_image,
-        strength=1-image_strength,
-        control_image=images,
-        num_inference_steps=20,
-        guidance_scale = guidance_scale,
-        controlnet_conditioning_scale=[face_strength, depth_control_scale],
-    ).images[0]
     
+    image = generate_image(prompt, negative, face_emb, face_image, image_strength, images, guidance_scale, face_strength, depth_control_scale, last_lora, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index)
     last_lora = repo_name
     return image, gr.update(visible=True)