Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 import torch
-import spaces
+import spaces  # Make sure this is imported
+import time
+from typing import Optional, List
+import numpy as np
+from PIL import Image
 torch.jit.script = lambda f: f
 import timm
-import time
 
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 from safetensors.torch import load_file
@@ -28,7 +31,6 @@ from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, UNet2DConditio
 import cv2
 import torch
 import numpy as np
-from PIL import Image
 
 from insightface.app import FaceAnalysis
 from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
@@ -159,50 +161,6 @@ last_lora = ""
 last_fused = False
 lora_archive = "/data"
 
-# Enhanced face detection with better face quality filtering
-def detect_faces(face_image, use_multiple_faces=False):
-    """
-    Detect faces in the image with quality filtering
-    Returns: list of face info dictionaries, or empty list if no faces
-    """
-    try:
-        face_info_list = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
-        
-        if not face_info_list or len(face_info_list) == 0:
-            print("No faces detected")
-            return []
-        
-        # Filter faces by quality score if available
-        filtered_faces = []
-        for face_info in face_info_list:
-            # Check if face has minimum quality
-            if 'det_score' in face_info and face_info['det_score'] > 0.5:
-                filtered_faces.append(face_info)
-            elif 'det_score' not in face_info:
-                filtered_faces.append(face_info)
-        
-        if not filtered_faces:
-            print("No high-quality faces detected")
-            return []
-        
-        # Sort faces by size (largest first)
-        filtered_faces = sorted(
-            filtered_faces, 
-            key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]),
-            reverse=True
-        )
-        
-        if use_multiple_faces:
-            print(f"Detected {len(filtered_faces)} high-quality faces")
-            return filtered_faces
-        else:
-            print(f"Using largest face (detected {len(filtered_faces)} total)")
-            return [filtered_faces[0]]
-            
-    except Exception as e:
-        print(f"Face detection error: {e}")
-        return []
-
 def process_face_embeddings_separately(face_info_list):
     """
     Process face embeddings separately for multi-face generation
@@ -300,42 +258,152 @@ def check_selected(selected_state, custom_lora):
     if not selected_state and not custom_lora:
         raise gr.Error("You must select a style")
 
-def resize_image_aspect_ratio(img, max_dim=1280):
-    width, height = img.size
-    aspect_ratio = width / height
+def shuffle_gallery(sdxl_loras):
+    random.shuffle(sdxl_loras)
+    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
 
-    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)
+def classify_gallery(sdxl_loras):
+    sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get("likes", 0), reverse=True)
+    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
 
-    new_width = (new_width // 8) * 8
-    new_height = (new_height // 8) * 8
+def swap_gallery(order, sdxl_loras):
+    if(order == "random"):
+        return shuffle_gallery(sdxl_loras)
+    else:
+        return classify_gallery(sdxl_loras)
+        
+def deselect():
+    return gr.Gallery(selected_index=None)
 
-    return img.resize((new_width, new_height), Image.LANCZOS)
+def get_huggingface_safetensors(link):
+    split_link = link.split("/")
+    if(len(split_link) == 2):
+        model_card = ModelCard.load(link)
+        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+        trigger_word = model_card.data.get("instance_prompt", "")
+        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+        fs = HfFileSystem()
+        try:
+            list_of_files = fs.ls(link, detail=False)
+            for file in list_of_files:
+                if(file.endswith(".safetensors")):
+                    safetensors_name = file.replace("/", "_")
+                    if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+                        fs.get_file(file, lpath=f"{lora_archive}/{safetensors_name}")
+                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                    image_elements = file.split("/")
+                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+        except:
+            gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+        return split_link[1], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
 
+def get_civitai_safetensors(link):
+    link_split = link.split("civitai.com/")
+    pattern = re.compile(r'models\/(\d+)')
+    regex_match = pattern.search(link_split[1])
+    if(regex_match):
+        civitai_model_id = regex_match.group(1)
+    else:
+        gr.Warning("No CivitAI model id found in your URL")
+        raise Exception("No CivitAI model id found in your URL")
+    model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token={os.getenv('CIVITAI_TOKEN')}"
+    x = requests.get(model_request_url)
+    if(x.status_code != 200):
+        raise Exception("Invalid CivitAI URL")
+    model_data = x.json()
+    
+    if(model_data["type"] != "LORA"):
+        gr.Warning("The model isn't tagged at CivitAI as a LoRA")
+        raise Exception("The model isn't tagged at CivitAI as a LoRA")
+    
+    model_link_download = None
+    image_url = None
+    trigger_word = ""
+    for model in model_data["modelVersions"]:
+        if(model["baseModel"] == "SDXL 1.0"):
+            model_link_download = f"{model['downloadUrl']}/?token={os.getenv('CIVITAI_TOKEN')}"
+            safetensors_name = model["files"][0]["name"]
+            if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+                safetensors_file_request = requests.get(model_link_download)
+                if(safetensors_file_request.status_code != 200):
+                    raise Exception("Invalid CivitAI download link")
+                with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
+                    file.write(safetensors_file_request.content)
+            trigger_word = model.get("trainedWords", [""])[0]
+            for image in model["images"]:
+                if(image["nsfwLevel"] == 1):
+                    image_url = image["url"]
+                    break
+            break
+    
+    if(not model_link_download):
+        gr.Warning("We couldn't find a SDXL LoRA on the model you've sent")
+        raise Exception("We couldn't find a SDXL LoRA on the model you've sent")
+    return model_data["name"], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
+
+def check_custom_model(link):
+    if(link.startswith("https://")):
+        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
+            link_split = link.split("huggingface.co/")
+            return get_huggingface_safetensors(link_split[1])
+        elif(link.startswith("https://civitai.com") or link.startswith("https://www.civitai.com")):
+            return get_civitai_safetensors(link)
+    else: 
+        return get_huggingface_safetensors(link)
+
+def load_custom_lora(link):
+    if(link):
+        try:
+            title, path, trigger_word, image = check_custom_model(link)
+            card = f'''
+            <div class="custom_lora_card">
+              <span>Loaded custom LoRA:</span>
+              <div class="card_internal">
+                <img src="{image}" />
+                <div>
+                    <h3>{title}</h3>
+                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
+                </div>
+              </div>
+            </div>
+            '''
+            return gr.update(visible=True), card, gr.update(visible=True), [path, trigger_word], gr.Gallery(selected_index=None), f"Custom: {path}"
+        except Exception as e:
+            gr.Warning("Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content")
+            return gr.update(visible=True), "Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), "", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+
+def remove_custom_lora():
+    return "", gr.update(visible=False), gr.update(visible=False), None
 
+@spaces.GPU(duration=120)
 def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, 
              guidance_scale, depth_control_scale, sdxl_loras, custom_lora, use_multiple_faces=False, 
              progress=gr.Progress(track_tqdm=True)):
     """
     Enhanced run_lora with improved face preservation and landscape mode
+    FIXED: Proper ZeroGPU decorator, no nested GPU calls
     """
     print("Custom LoRA:", custom_lora)
     custom_lora_path = custom_lora[0] if custom_lora else None
     selected_state_index = selected_state.index if selected_state else -1
 
     st = time.time()
+    
+    # Ensure models are on GPU
+    pipe.to(device)
+    zoe.to(device)
+    
     face_image = resize_image_aspect_ratio(face_image)
 
-    # Enhanced face detection
+    # Enhanced face detection (CPU operation - InsightFace uses CPU)
     face_info_list = detect_faces(face_image, use_multiple_faces)
     face_detected = len(face_info_list) > 0
     
     if face_detected:
-        # CHANGED: Process faces separately instead of averaging
+        # Process faces separately instead of averaging
         face_embeddings = process_face_embeddings_separately(face_info_list)
         face_kps = create_face_kps_image(face_image, face_info_list)
         print(f"Processing with {len(face_info_list)} face(s) separately")
@@ -396,7 +464,7 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     et = time.time()
     print('Prompt processing took:', et - st, 'seconds')
 
-    # IMPROVED: Better parameter adjustment for face/landscape modes
+    # Better parameter adjustment for face/landscape modes
     if not face_detected:
         # Enhanced landscape mode parameters
         face_strength = 0.0
@@ -411,22 +479,32 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
 
     st = time.time()
     
-    # Generate single image with best face (or landscape)
-    image = generate_image(
-        prompt, negative, face_emb, face_image, face_kps, image_strength,
-        guidance_scale, face_strength, depth_control_scale, repo_name,
-        full_path_lora, lora_scale, sdxl_loras, selected_state_index, face_detected, st
-    )
-
+    # FIXED: Call non-decorated version (inline generation)
+    try:
+        image = generate_image_inline(
+            prompt, negative, face_emb, face_image, face_kps, image_strength,
+            guidance_scale, face_strength, depth_control_scale, repo_name,
+            full_path_lora, lora_scale, sdxl_loras, selected_state_index, face_detected, st
+        )
+    except Exception as e:
+        print(f"Generation error: {e}")
+        torch.cuda.empty_cache()
+        raise gr.Error(f"Image generation failed: {str(e)}")
+    
+    # Cleanup GPU memory
+    torch.cuda.empty_cache()
+    
     return (face_image, image), gr.update(visible=True)
 
-run_lora.zerogpu = True
-
 
-@spaces.GPU(duration=90)  # Increased duration for better quality
-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, st):
+# FIXED: Removed @spaces.GPU decorator - this runs within GPU context
+def generate_image_inline(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, st):
+    """
+    FIXED: No decorator - called from within GPU context
+    """
     global last_fused, last_lora
 
     print("Loaded state dict:", loaded_state_dict)
@@ -502,6 +580,8 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     num_inference_steps = 50  # Increased for better quality
     
     print("Generating image...")
+    print(f"GPU Memory before generation: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
+    
     image = pipe(
         prompt_embeds=conditioning,
         pooled_prompt_embeds=pooled,
@@ -518,128 +598,79 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
         controlnet_conditioning_scale=control_scales,
     ).images[0]
 
+    print(f"GPU Memory after generation: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
+    
     last_lora = repo_name
     return image
 
-def shuffle_gallery(sdxl_loras):
-    random.shuffle(sdxl_loras)
-    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
 
-def classify_gallery(sdxl_loras):
-    sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get("likes", 0), reverse=True)
-    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
-
-def swap_gallery(order, sdxl_loras):
-    if(order == "random"):
-        return shuffle_gallery(sdxl_loras)
-    else:
-        return classify_gallery(sdxl_loras)
+# CPU-bound helper functions (no decorators needed)
+def detect_faces(face_image, use_multiple_faces=False):
+    """
+    Detect faces in the image with quality filtering
+    CPU operation - no GPU decorator needed
+    """
+    try:
+        face_info_list = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
         
-def deselect():
-    return gr.Gallery(selected_index=None)
+        if not face_info_list or len(face_info_list) == 0:
+            print("No faces detected")
+            return []
+        
+        # Filter faces by quality score if available
+        filtered_faces = []
+        for face_info in face_info_list:
+            # Check if face has minimum quality
+            if 'det_score' in face_info and face_info['det_score'] > 0.5:
+                filtered_faces.append(face_info)
+            elif 'det_score' not in face_info:
+                filtered_faces.append(face_info)
+        
+        if not filtered_faces:
+            print("No high-quality faces detected")
+            return []
+        
+        # Sort faces by size (largest first)
+        filtered_faces = sorted(
+            filtered_faces, 
+            key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]),
+            reverse=True
+        )
+        
+        if use_multiple_faces:
+            print(f"Detected {len(filtered_faces)} high-quality faces")
+            return filtered_faces
+        else:
+            print(f"Using largest face (detected {len(filtered_faces)} total)")
+            return [filtered_faces[0]]
+            
+    except Exception as e:
+        print(f"Face detection error: {e}")
+        return []
 
-def get_huggingface_safetensors(link):
-    split_link = link.split("/")
-    if(len(split_link) == 2):
-        model_card = ModelCard.load(link)
-        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-        trigger_word = model_card.data.get("instance_prompt", "")
-        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-        fs = HfFileSystem()
-        try:
-            list_of_files = fs.ls(link, detail=False)
-            for file in list_of_files:
-                if(file.endswith(".safetensors")):
-                    safetensors_name = file.replace("/", "_")
-                    if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
-                        fs.get_file(file, lpath=f"{lora_archive}/{safetensors_name}")
-                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                    image_elements = file.split("/")
-                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-        except:
-            gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-        return split_link[1], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
 
-def get_civitai_safetensors(link):
-    link_split = link.split("civitai.com/")
-    pattern = re.compile(r'models\/(\d+)')
-    regex_match = pattern.search(link_split[1])
-    if(regex_match):
-        civitai_model_id = regex_match.group(1)
-    else:
-        gr.Warning("No CivitAI model id found in your URL")
-        raise Exception("No CivitAI model id found in your URL")
-    model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token={os.getenv('CIVITAI_TOKEN')}"
-    x = requests.get(model_request_url)
-    if(x.status_code != 200):
-        raise Exception("Invalid CivitAI URL")
-    model_data = x.json()
-    
-    if(model_data["type"] != "LORA"):
-        gr.Warning("The model isn't tagged at CivitAI as a LoRA")
-        raise Exception("The model isn't tagged at CivitAI as a LoRA")
-    
-    model_link_download = None
-    image_url = None
-    trigger_word = ""
-    for model in model_data["modelVersions"]:
-        if(model["baseModel"] == "SDXL 1.0"):
-            model_link_download = f"{model['downloadUrl']}/?token={os.getenv('CIVITAI_TOKEN')}"
-            safetensors_name = model["files"][0]["name"]
-            if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
-                safetensors_file_request = requests.get(model_link_download)
-                if(safetensors_file_request.status_code != 200):
-                    raise Exception("Invalid CivitAI download link")
-                with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
-                    file.write(safetensors_file_request.content)
-            trigger_word = model.get("trainedWords", [""])[0]
-            for image in model["images"]:
-                if(image["nsfwLevel"] == 1):
-                    image_url = image["url"]
-                    break
-            break
-    
-    if(not model_link_download):
-        gr.Warning("We couldn't find a SDXL LoRA on the model you've sent")
-        raise Exception("We couldn't find a SDXL LoRA on the model you've sent")
-    return model_data["name"], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
+def resize_image_aspect_ratio(img, max_dim=1280):
+    """CPU operation"""
+    width, height = img.size
+    aspect_ratio = width / height
 
-def check_custom_model(link):
-    if(link.startswith("https://")):
-        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
-            link_split = link.split("huggingface.co/")
-            return get_huggingface_safetensors(link_split[1])
-        elif(link.startswith("https://civitai.com") or link.startswith("https://www.civitai.com")):
-            return get_civitai_safetensors(link)
-    else: 
-        return get_huggingface_safetensors(link)
+    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)
 
-def load_custom_lora(link):
-    if(link):
-        try:
-            title, path, trigger_word, image = check_custom_model(link)
-            card = f'''
-            <div class="custom_lora_card">
-              <span>Loaded custom LoRA:</span>
-              <div class="card_internal">
-                <img src="{image}" />
-                <div>
-                    <h3>{title}</h3>
-                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
-                </div>
-              </div>
-            </div>
-            '''
-            return gr.update(visible=True), card, gr.update(visible=True), [path, trigger_word], gr.Gallery(selected_index=None), f"Custom: {path}"
-        except Exception as e:
-            gr.Warning("Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content")
-            return gr.update(visible=True), "Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
-    else:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
+
+    return img.resize((new_width, new_height), Image.LANCZOS)
 
-def remove_custom_lora():
-    return "", gr.update(visible=False), gr.update(visible=False), None
+
+def check_selected(selected_state, custom_lora):
+    """CPU operation"""
+    if not selected_state and not custom_lora:
+        raise gr.Error("You must select a style")
 
 # Build Gradio interface
 with gr.Blocks(css="custom.css") as demo:
@@ -755,7 +786,7 @@ with gr.Blocks(css="custom.css") as demo:
         inputs=[selected_state, custom_loaded_lora],
         show_progress=False
     ).success(
-        fn=run_lora,
+        fn=run_lora,  # This now has proper @spaces.GPU decorator
         inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, 
                guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora, use_multiple_faces],
         outputs=[result, share_group],