Update app.py

app.py

@@ -173,23 +173,6 @@ button.addEventListener('click', function() {
 '''
 lora_archive = "/data"
 
-def resize_image_aspect_ratio(img, max_dim=512):
-    width, height = img.size
-    aspect_ratio = width / height
-
-    if aspect_ratio >= 1:  # Landscape or square
-        new_width = min(max_dim, width)
-        new_height = int(new_width / aspect_ratio)
-    else:  # Portrait
-        new_height = min(max_dim, height)
-        new_width = int(new_height * aspect_ratio)
-
-    new_width = (new_width // 8) * 8
-    new_height = (new_height // 8) * 8
-
-    return img.resize((new_width, new_height), Image.LANCZOS)
-
-
 def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative, is_new=False):
     lora_repo = sdxl_loras[selected_state.index]["repo"]
     new_placeholder = "Type a prompt to use your selected LoRA"
@@ -246,57 +229,22 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 )
                 del weights_sd
                 del lora_model
+def resize_image_aspect_ratio(img, max_dim=512):
+    width, height = img.size
+    aspect_ratio = width / height
 
-@spaces.GPU(duration=80)
-def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, face_detected):
-    global last_fused, last_lora
-
-    control_images = [face_kps, zoe(face_image)] if face_detected else [zoe(face_image)]
-    control_scales = [face_strength, depth_control_scale] if face_detected else [depth_control_scale]
-
-    if repo_name.startswith("https://huggingface.co"):
-      repo_id = repo_name.split("huggingface.co/")[-1]
-      weight_file = "pytorch_lora_weights.safetensors"
-      full_path_lora = hf_hub_download(repo_id=repo_id, filename=weight_file, repo_type="model")
-      loaded_state_dict = load_file(full_path_lora)
-
-    if last_lora != repo_name:
-        if last_fused:
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
-            pipe.unload_textual_inversion()
-        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:
-            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)
-            pipe.load_textual_inversion(state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
-            pipe.load_textual_inversion(state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+    if aspect_ratio >= 1:  # Landscape or square
+        new_width = min(max_dim, width)
+        new_height = int(new_width / aspect_ratio)
+    else:  # Portrait
+        new_height = min(max_dim, height)
+        new_width = int(new_height * aspect_ratio)
 
-    conditioning, pooled = compel(prompt)
-    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
 
-    image = pipe(
-        prompt_embeds=conditioning,
-        pooled_prompt_embeds=pooled,
-        negative_prompt_embeds=negative_conditioning,
-        negative_pooled_prompt_embeds=negative_pooled,
-        width=face_image.width,
-        height=face_image.height,
-        image_embeds=face_emb if face_detected else None,
-        image=face_image,
-        strength=1-image_strength,
-        control_image=control_images,
-        num_inference_steps=36,
-        guidance_scale=guidance_scale,
-        controlnet_conditioning_scale=control_scales,
-    ).images[0]
+    return img.resize((new_width, new_height), Image.LANCZOS)
 
-    last_lora = repo_name
-    return image
 
 
 def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, custom_lora, progress=gr.Progress(track_tqdm=True)):
@@ -350,8 +298,77 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
 
     return (resized_image, image), gr.update(visible=True)
 
+@spaces.GPU(duration=100)
+def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, face_detected):
+    global last_fused, last_lora
+
+    print("Loaded state dict:", loaded_state_dict)
+    print("Last LoRA:", last_lora, "| Current LoRA:", repo_name)
+
+    control_images = [face_kps, zoe(face_image)] if face_detected else [zoe(face_image)]
+    control_scales = [face_strength, depth_control_scale] if face_detected else [depth_control_scale]
+
+    # Handle Hugging Face URL-based LoRA
+    if repo_name.startswith("https://huggingface.co"):
+        repo_id = repo_name.split("huggingface.co/")[-1]
+        fs = HfFileSystem()
+        files = fs.ls(repo_id, detail=False)
+        safetensors_files = [f for f in files if f.endswith(".safetensors")]
+
+        if not safetensors_files:
+            raise gr.Error("No .safetensors file found in this Hugging Face repository.")
+
+        weight_file = safetensors_files[0]  # Dynamically select the first available .safetensors file
+        full_path_lora = hf_hub_download(repo_id=repo_id, filename=weight_file, repo_type="model")
+        loaded_state_dict = load_file(full_path_lora)
+    else:
+        # Use the previously loaded state_dict if not using a Hugging Face URL
+        loaded_state_dict = state_dicts[repo_name]["state_dict"]
+
+    # Manage LoRA weights and textual inversion embeddings
+    if last_lora != repo_name:
+        if last_fused:
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            pipe.unload_textual_inversion()
+        pipe.load_lora_weights(loaded_state_dict)
+        pipe.fuse_lora(lora_scale)
+        last_fused = True
+
+        # Handle pivotal tuning (textual inversion embeddings)
+        is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
+        if is_pivotal:
+            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)
+            pipe.load_textual_inversion(state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+            pipe.load_textual_inversion(state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+
+    # Prompt embeddings
+    print("Processing prompt...")
+    conditioning, pooled = compel(prompt)
+    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
+
+    # Image generation
+    print("Generating image...")
+    image = pipe(
+        prompt_embeds=conditioning,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=negative_conditioning,
+        negative_pooled_prompt_embeds=negative_pooled,
+        width=face_image.width,
+        height=face_image.height,
+        image_embeds=face_emb if face_detected else None,
+        image=face_image,
+        strength=1-image_strength,
+        control_image=control_images,
+        num_inference_steps=36,
+        guidance_scale=guidance_scale,
+        controlnet_conditioning_scale=control_scales,
+    ).images[0]
 
-run_lora.zerogpu = True
+    last_lora = repo_name
+    return image
 
 def shuffle_gallery(sdxl_loras):
     random.shuffle(sdxl_loras)