Update app.py

app.py

@@ -1,26 +1,9 @@
-import os
-# Fix transformers deprecation warning - use HF_HOME instead of TRANSFORMERS_CACHE
-if 'TRANSFORMERS_CACHE' in os.environ and 'HF_HOME' not in os.environ:
-    os.environ['HF_HOME'] = os.environ['TRANSFORMERS_CACHE']
-if 'HF_HOME' not in os.environ:
-    os.environ['HF_HOME'] = os.environ.get('TRANSFORMERS_CACHE', '/data/.cache/huggingface')
-
-# Suppress ONNX Runtime GPU discovery warnings (we use CPU for face detection)
-os.environ['ORT_LOGGING_LEVEL'] = '3'  # Only show errors, not warnings
-
-import warnings
-warnings.filterwarnings('ignore', category=UserWarning, module='onnxruntime')
-
 import gradio as gr
 import torch
 import spaces
-import time
-import traceback
-from typing import Optional, List
-import numpy as np
-from PIL import Image, ImageEnhance
 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
@@ -45,6 +28,7 @@ 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
@@ -80,248 +64,174 @@ with open("defaults_data.json", "r") as file:
 
 device = "cuda"
 
-print("=" * 80)
-print("🚀 Starting LucasArts Style App (VRAM Optimized)")
-print("=" * 80)
-print()
-
 # Cache for LoRA state dicts
-print("📦 Loading LoRA configurations...")
 state_dicts = {}
-try:
-    for item in sdxl_loras_raw:
-        saved_name = hf_hub_download(item["repo"], item["weights"])
-        
-        if not saved_name.endswith('.safetensors'):
-            state_dict = torch.load(saved_name)
-        else:
-            state_dict = load_file(saved_name)
-        
-        state_dicts[item["repo"]] = {
-            "saved_name": saved_name,
-            "state_dict": state_dict
-        }
-    print(f"✅ Loaded {len(state_dicts)} LoRA configurations")
-except Exception as e:
-    print(f"❌ Error loading LoRAs: {e}")
-    raise
+for item in sdxl_loras_raw:
+    saved_name = hf_hub_download(item["repo"], item["weights"])
+    
+    if not saved_name.endswith('.safetensors'):
+        state_dict = torch.load(saved_name)
+    else:
+        state_dict = load_file(saved_name)
+    
+    state_dicts[item["repo"]] = {
+        "saved_name": saved_name,
+        "state_dict": state_dict
+    }
 
 sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
 
 # Download models
-print()
-print("📥 Downloading required models...")
-try:
-    hf_hub_download(
-        repo_id="InstantX/InstantID",
-        filename="ControlNetModel/config.json",
-        local_dir="/data/checkpoints",
-    )
-    hf_hub_download(
-        repo_id="InstantX/InstantID",
-        filename="ControlNetModel/diffusion_pytorch_model.safetensors",
-        local_dir="/data/checkpoints",
-    )
-    hf_hub_download(
-        repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="/data/checkpoints"
-    )
-    hf_hub_download(
-        repo_id="latent-consistency/lcm-lora-sdxl",
-        filename="pytorch_lora_weights.safetensors",
-        local_dir="/data/checkpoints",
-    )
-    print("✅ Model checkpoints downloaded")
-except Exception as e:
-    print(f"❌ Error downloading models: {e}")
-    raise
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/config.json",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/diffusion_pytorch_model.safetensors",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="/data/checkpoints"
+)
+hf_hub_download(
+    repo_id="latent-consistency/lcm-lora-sdxl",
+    filename="pytorch_lora_weights.safetensors",
+    local_dir="/data/checkpoints",
+)
 
 # Download antelopev2
-print()
-print("📥 Downloading face detection model...")
-try:
-    antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
-    print(f"✅ Face detection model: {antelope_download}")
-except Exception as e:
-    print(f"❌ Error downloading face model: {e}")
-    raise
-
-print()
-print("🔧 Initializing face detection...")
-# VRAM OPTIMIZED: Standard 768x768 for better memory usage
-try:
-    app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
-    app.prepare(ctx_id=0, det_size=(768, 768))
-    print("✅ Face detection initialized at 768x768")
-except Exception as e:
-    print(f"❌ Error initializing face detection: {e}")
-    raise
+antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
+print(antelope_download)
+app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(768, 768))
 
 # Prepare models
 face_adapter = f'/data/checkpoints/ip-adapter.bin'
 controlnet_path = f'/data/checkpoints/ControlNetModel'
 
-print()
-print("🔧 Loading ControlNets...")
 st = time.time()
-try:
-    identitynet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
-    zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0", torch_dtype=torch.float16)
-    et = time.time()
-    print(f'✅ ControlNet loaded in {et - st:.2f} seconds')
-except Exception as e:
-    print(f"❌ Error loading ControlNet: {e}")
-    raise
+identitynet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
+zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0", torch_dtype=torch.float16)
+et = time.time()
+print('Loading ControlNet took: ', et - st, 'seconds')
 
-print()
-print("🔧 Loading VAE...")
 st = time.time()
-try:
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-    et = time.time()
-    print(f'✅ VAE loaded in {et - st:.2f} seconds')
-except Exception as e:
-    print(f"❌ Error loading VAE: {e}")
-    raise
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+et = time.time()
+print('Loading VAE took: ', et - st, 'seconds')
 
-print()
-print("🔧 Loading main pipeline...")
 st = time.time()
-try:
-    pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
-        "frankjoshua/albedobaseXL_v21", 
-        vae=vae,
-        controlnet=[identitynet, zoedepthnet],
-        torch_dtype=torch.float16
-    )
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
+    "SG161222/RealVisXL_V5.0", 
+    vae=vae,
+    controlnet=[identitynet, zoedepthnet],
+    torch_dtype=torch.float16
+)
+
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
+pipe.load_ip_adapter_instantid(face_adapter)
+pipe.set_ip_adapter_scale(0.9)
+et = time.time()
+print('Loading pipeline took: ', et - st, 'seconds')
 
-    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
-    pipe.load_ip_adapter_instantid(face_adapter)
-    pipe.set_ip_adapter_scale(1.0)
-    et = time.time()
-    print(f'✅ Pipeline loaded in {et - st:.2f} seconds')
-except Exception as e:
-    print(f"❌ Error loading pipeline: {e}")
-    raise
-
-print()
-print("🔧 Loading Compel (prompt processor)...")
 st = time.time()
-try:
-    compel = Compel(
-        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-        requires_pooled=[False, True]
-    )
-    et = time.time()
-    print(f'✅ Compel loaded in {et - st:.2f} seconds')
-except Exception as e:
-    print(f"❌ Error loading Compel: {e}")
-    raise
+compel = Compel(
+    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+    requires_pooled=[False, True]
+)
+et = time.time()
+print('Loading Compel took: ', et - st, 'seconds')
 
-print()
-print("🔧 Loading Zoe (depth detector)...")
 st = time.time()
-try:
-    zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-    et = time.time()
-    print(f'✅ Zoe loaded in {et - st:.2f} seconds')
-except Exception as e:
-    print(f"❌ Error loading Zoe: {e}")
-    raise
-
-print()
-print("🔧 Moving models to GPU...")
-try:
-    zoe.to(device)
-    pipe.to(device)
-    print("✅ Models moved to GPU")
-except Exception as e:
-    print(f"❌ Error moving to GPU: {e}")
-    raise
-
-print()
-print("=" * 80)
-print("✅ All models loaded successfully!")
-print("âš¡ VRAM Optimized Configuration:")
-print("   • 768x768 face detection")
-print("   • 768px max output resolution")
-print("   • 35 inference steps")
-print("   • Enhanced error reporting")
-print("   • Fixed version compatibility (diffusers 0.21.4)")
-print("=" * 80)
-print()
+zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
+et = time.time()
+print('Loading Zoe took: ', et - st, 'seconds')
+zoe.to(device)
+pipe.to(device)
 
 last_lora = ""
 last_fused = False
 lora_archive = "/data"
 
-
-def enhance_details(image, strength=1.15):
-    """Post-process to enhance details"""
+# Improved face detection with multi-face support
+def detect_faces(face_image, use_multiple_faces=False):
+    """
+    Detect faces in the image
+    Returns: list of face info dictionaries, or empty list if no faces
+    """
     try:
-        sharpener = ImageEnhance.Sharpness(image)
-        image = sharpener.enhance(strength)
-        
-        contrast = ImageEnhance.Contrast(image)
-        image = contrast.enhance(1.08)
+        face_info_list = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
         
-        return image
-    except Exception as e:
-        print(f"Warning: Detail enhancement failed: {e}")
-        return image
-
-
-def enhanced_depth_map(image, face_detected=False):
-    """Better depth map generation"""
-    try:
-        original_size = image.size
+        if not face_info_list or len(face_info_list) == 0:
+            print("No faces detected")
+            return []
         
-        # Skip upscaling for VRAM optimization
-        depth = zoe(image)
+        # Sort faces by size (largest first)
+        face_info_list = sorted(
+            face_info_list, 
+            key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]),
+            reverse=True
+        )
         
-        return depth
+        if use_multiple_faces:
+            print(f"Detected {len(face_info_list)} faces")
+            return face_info_list
+        else:
+            print(f"Using largest face (detected {len(face_info_list)} total)")
+            return [face_info_list[0]]
+            
     except Exception as e:
-        print(f"Error in depth map generation: {e}")
-        # Return a blank depth map as fallback
-        return Image.new('L', image.size, color=128)
-
+        print(f"Face detection error: {e}")
+        return []
 
-def process_face_embeddings_separately(face_info_list):
-    """Process face embeddings separately for multi-face generation"""
+def process_face_embeddings(face_info_list):
+    """
+    Process face embeddings - average multiple faces or return single face
+    """
     if not face_info_list:
-        return []
+        return None
     
+    if len(face_info_list) == 1:
+        return face_info_list[0]['embedding']
+    
+    # Average embeddings for multiple faces
     embeddings = [face_info['embedding'] for face_info in face_info_list]
-    return embeddings
-
+    avg_embedding = np.mean(embeddings, axis=0)
+    return avg_embedding
 
 def create_face_kps_image(face_image, face_info_list):
-    """Create keypoints image from face info with enhanced visibility"""
+    """
+    Create keypoints image from face info
+    """
     if not face_info_list:
         return face_image
     
+    # For multiple faces, draw all keypoints
     if len(face_info_list) > 1:
         return draw_multiple_kps(face_image, [f['kps'] for f in face_info_list])
     else:
         return draw_kps(face_image, face_info_list[0]['kps'])
 
-
 def draw_multiple_kps(image_pil, kps_list, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]):
-    """Draw keypoints for multiple faces with enhanced visibility"""
+    """
+    Draw keypoints for multiple faces
+    """
     stickwidth = 4
     limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
     
     w, h = image_pil.size
     out_img = np.zeros([h, w, 3])
     
-    for idx, kps in enumerate(kps_list):
+    for kps in kps_list:
         kps = np.array(kps)
-        color_offset = idx % len(color_list)
         
         for i in range(len(limbSeq)):
             index = limbSeq[i]
-            color = color_list[(index[0] + color_offset) % len(color_list)]
+            color = color_list[index[0]]
             
             x = kps[index][:, 0]
             y = kps[index][:, 1]
@@ -335,25 +245,24 @@ def draw_multiple_kps(image_pil, kps_list, color_list=[(255, 0, 0), (0, 255, 0),
         out_img = (out_img * 0.6).astype(np.uint8)
         
         for idx_kp, kp in enumerate(kps):
-            color = color_list[(idx_kp + color_offset) % len(color_list)]
+            color = color_list[idx_kp]
             x, y = kp
             out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
     
     out_img_pil = Image.fromarray(out_img.astype(np.uint8))
     return out_img_pil
 
-
 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"
     weight_name = sdxl_loras[selected_state.index]["weights"]
-    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
 
     for lora_list in lora_defaults:
         if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
-            face_strength = lora_list.get("face_strength", 1.0)
-            image_strength = lora_list.get("image_strength", 0.15)
-            weight = lora_list.get("weight", 1.0)
+            face_strength = lora_list.get("face_strength", 0.9)
+            image_strength = lora_list.get("image_strength", 0.2)
+            weight = lora_list.get("weight", 0.95)
             depth_control_scale = lora_list.get("depth_control_scale", 0.8)
             negative = lora_list.get("negative", "")
     
@@ -374,33 +283,230 @@ def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, i
         selected_state
     )
 
-
 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
+
+    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 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 support for:
+    - No faces (landscape mode)
+    - Multiple faces
+    - Improved results
+    """
+    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()
+    face_image = resize_image_aspect_ratio(face_image)
+
+    # Enhanced face detection
+    face_info_list = detect_faces(face_image, use_multiple_faces)
+    face_detected = len(face_info_list) > 0
+    
+    if face_detected:
+        face_emb = process_face_embeddings(face_info_list)
+        face_kps = create_face_kps_image(face_image, face_info_list)
+        print(f"Processing with {len(face_info_list)} face(s)")
+    else:
+        face_emb = None
+        face_kps = face_image
+        print("No faces detected - using landscape/depth mode only")
+
+    et = time.time()
+    print('Face processing took:', et - st, 'seconds')
+
+    st = time.time()
+
+    # Enhanced prompt processing
+    if custom_lora_path and custom_lora[1]:
+        prompt = f"{prompt} {custom_lora[1]}"
+    else:
+        for lora_list in lora_defaults:
+            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)
+
+    print("Prompt:", prompt)
+    if prompt == "":
+        prompt = "a beautiful scene" if not face_detected else "a person"
+    print(f"Executing prompt: {prompt}")
+
+    if negative == "":
+        # Enhanced negative prompt for better quality
+        negative = "worst quality, low quality, blurry, distorted, deformed" if not face_detected else None
+    
+    print("Custom Loaded LoRA:", custom_lora_path)
+
+    if not selected_state and not custom_lora_path:
+        raise gr.Error("You must select a style")
+    elif custom_lora_path:
+        repo_name = custom_lora_path
+        full_path_lora = custom_lora_path
+    else:
+        repo_name = sdxl_loras[selected_state_index]["repo"]
+        full_path_lora = state_dicts[repo_name]["saved_name"]
+
+    repo_name = repo_name.rstrip("/").lower()
+
+    print("Full path LoRA", full_path_lora)
+
+    et = time.time()
+    print('Prompt processing took:', et - st, 'seconds')
+
+    # Adjust parameters based on face detection
+    if not face_detected:
+        # For landscape/no face mode, reduce face strength and increase depth control
+        face_strength = 0.0
+        depth_control_scale = max(depth_control_scale, 0.9)
+        image_strength = min(image_strength, 0.4)
+        print("Adjusted parameters for no-face mode")
+
+    st = time.time()
+    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
+    )
+
+    return (face_image, image), gr.update(visible=True)
+
+run_lora.zerogpu = True
+
+
+@spaces.GPU(duration=75)
+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):
+    global last_fused, last_lora
+
+    print("Loaded state dict:", loaded_state_dict)
+    print("Last LoRA:", last_lora, "| Current LoRA:", repo_name)
+
+    # Prepare control images and scales based on face detection
+    if face_detected:
+        control_images = [face_kps, zoe(face_image)]
+        control_scales = [face_strength, depth_control_scale]
+    else:
+        # Only use depth control for landscapes
+        control_images = [zoe(face_image)]
+        control_scales = [depth_control_scale]
+
+    # Handle custom LoRA from HuggingFace
+    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]
+        full_path_lora = hf_hub_download(repo_id=repo_id, filename=weight_file, repo_type="model")
+    else:
+        full_path_lora = loaded_state_dict
+
+    # Improved LoRA loading and caching
+    if last_lora != repo_name:
+        if last_fused:
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            pipe.unload_textual_inversion()
+
+        # Load LoRA with better error handling
+        try:
+            pipe.load_lora_weights(full_path_lora)
+            pipe.fuse_lora(lora_scale)
+            last_fused = True
+            
+            # Handle pivotal tuning 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
+                )
+        except Exception as e:
+            print(f"Error loading LoRA: {e}")
+            raise gr.Error(f"Failed to load LoRA: {str(e)}")
+
+    print("Processing prompt...")
+    conditioning, pooled = compel(prompt)
+    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
+
+    # Enhanced generation parameters
+    num_inference_steps = 40  # Increased for better quality
+    
+    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=num_inference_steps,
+        guidance_scale=guidance_scale,
+        controlnet_conditioning_scale=control_scales,
+    ).images[0]
+
+    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)
         
-
 def deselect():
     return gr.Gallery(selected_index=None)
 
-
 def get_huggingface_safetensors(link):
     split_link = link.split("/")
     if(len(split_link) == 2):
@@ -424,7 +530,6 @@ def get_huggingface_safetensors(link):
             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+)')
@@ -469,7 +574,6 @@ def get_civitai_safetensors(link):
         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")):
@@ -480,7 +584,6 @@ def check_custom_model(link):
     else: 
         return get_huggingface_safetensors(link)
 
-
 def load_custom_lora(link):
     if(link):
         try:
@@ -504,419 +607,20 @@ def load_custom_lora(link):
     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)):
-    """
-    VRAM optimized generation with enhanced error reporting
-    """
-    print("=" * 80)
-    print("🚀 FUNCTION CALLED: run_lora")
-    print(f"📊 Inputs received:")
-    print(f"  - face_image: {type(face_image)} - {face_image.size if face_image else 'None'}")
-    print(f"  - prompt: '{prompt}'")
-    print(f"  - selected_state: {type(selected_state)} - {selected_state}")
-    print(f"  - custom_lora: {custom_lora}")
-    print(f"  - use_multiple_faces: {use_multiple_faces}")
-    print("=" * 80)
-    
-    try:
-        print("Starting generation...")
-        print("Custom LoRA:", custom_lora)
-        custom_lora_path = custom_lora[0] if custom_lora else None
-        
-        # Extract index from selected_state (handle Gradio SelectData object)
-        if selected_state:
-            print(f"  selected_state exists: {selected_state}")
-            print(f"  selected_state type: {type(selected_state)}")
-            print(f"  selected_state dir: {dir(selected_state)}")
-            if hasattr(selected_state, 'index'):
-                selected_state_index = selected_state.index
-                print(f"  ✓ Extracted index: {selected_state_index}")
-            else:
-                selected_state_index = -1
-                print(f"  ❌ No index attribute, using -1")
-        else:
-            selected_state_index = -1
-            print(f"  ❌ selected_state is None or False")
-        
-        print(f"🔍 VALIDATION CHECK:")
-        print(f"  - selected_state_index: {selected_state_index}")
-        print(f"  - custom_lora_path: {custom_lora_path}")
-        print(f"  - len(sdxl_loras): {len(sdxl_loras)}")
-
-        # Validate selection immediately
-        if (selected_state_index is None or selected_state_index < 0) and not custom_lora_path:
-            error_msg = "❌ You must select a style before generating"
-            print(error_msg)
-            return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-        
-        # Validate selected_state_index is valid (only check positive indices)
-        if not custom_lora_path and selected_state_index >= 0 and selected_state_index >= len(sdxl_loras):
-            error_msg = f"❌ Invalid style selection (index: {selected_state_index}, available: {len(sdxl_loras)})"
-            print(error_msg)
-            return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-
-        st = time.time()
-        
-        pipe.to(device)
-        zoe.to(device)
-        
-        # VRAM OPTIMIZED: Reduced max dimension to 768
-        face_image = resize_image_aspect_ratio(face_image)
-        print(f"Resized image to {face_image.size}")
-
-        # Face detection with better error handling
-        try:
-            face_info_list = detect_faces(face_image, use_multiple_faces)
-            face_detected = len(face_info_list) > 0
-        except Exception as e:
-            print(f"Face detection error: {e}")
-            face_detected = False
-            face_info_list = []
-        
-        if face_detected:
-            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) detected")
-            face_emb = face_embeddings[0]
-        else:
-            face_emb = None
-            face_kps = face_image
-            print("No faces detected - landscape mode")
-
-        et = time.time()
-        print(f'Face processing took: {et - st:.2f}s')
-
-        st = time.time()
-
-        # Enhanced prompt processing
-        if custom_lora_path and custom_lora[1]:
-            prompt = f"{prompt} {custom_lora[1]}"
-        else:
-            for lora_list in lora_defaults:
-                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)
-
-        if "lucasarts artstyle" not in prompt.lower():
-            prompt = f"{prompt}, lucasarts artstyle"
-
-        print("Prompt:", prompt)
-        if prompt == "":
-            prompt = "a beautiful cinematic scene" if not face_detected else "a person in cinematic lighting"
-        print(f"Executing prompt: {prompt}")
-
-        if negative == "":
-            if not face_detected:
-                negative = "worst quality, low quality, blurry, distorted, deformed, ugly, bad anatomy"
-            else:
-                negative = "worst quality, low quality, blurry, distorted, deformed, ugly, bad anatomy, bad proportions"
-        
-        print("Custom Loaded LoRA:", custom_lora_path)
-
-        if custom_lora_path:
-            repo_name = custom_lora_path
-            full_path_lora = custom_lora_path
-        else:
-            repo_name = sdxl_loras[selected_state_index]["repo"]
-            if repo_name not in state_dicts:
-                error_msg = f"❌ LoRA not loaded: {repo_name}\nAvailable: {list(state_dicts.keys())[:5]}"
-                print(error_msg)
-                return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-            full_path_lora = state_dicts[repo_name]["saved_name"]
-
-        repo_name = repo_name.rstrip("/").lower()
-
-        et = time.time()
-        print(f'Prompt processing took: {et - st:.2f}s')
-
-        # Optimized parameters based on mode
-        if not face_detected:
-            face_strength = 0.0
-            depth_control_scale = 1.0
-            image_strength = 0.25
-            guidance_scale = max(guidance_scale, 8.5)
-            print("Optimized for landscape mode")
-        else:
-            face_strength = max(face_strength, 1.0)
-            depth_control_scale = max(depth_control_scale, 0.8)
-            guidance_scale = max(guidance_scale, 7.5)
-            print("Optimized for face preservation")
-
-        st = time.time()
-        
-        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
-        )
-        
-        torch.cuda.empty_cache()
-        
-        print("Generation complete!")
-        print("=" * 80)
-        return (face_image, image), gr.update(visible=True), gr.update(visible=False)
-        
-    except torch.cuda.OutOfMemoryError as e:
-        error_msg = (
-            "GPU OUT OF MEMORY!\n\n"
-            "Your image is too large for available VRAM.\n\n"
-            "Solutions:\n"
-            "1. Try a smaller image (current max: 768px)\n"
-            "2. Upgrade to A10G GPU in Space settings\n"
-            "3. Reduce image strength parameter\n\n"
-            f"Technical details: {str(e)}"
-        )
-        print("=" * 80)
-        print(error_msg)
-        print("=" * 80)
-        torch.cuda.empty_cache()
-        return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-    except RuntimeError as e:
-        if "out of memory" in str(e).lower():
-            error_msg = (
-                "GPU OUT OF MEMORY!\n\n"
-                f"Error: {str(e)}\n\n"
-                "Solutions:\n"
-                "1. Upload smaller image\n"
-                "2. Upgrade GPU in Settings\n"
-                "3. Reduce parameters"
-            )
-        else:
-            error_msg = f"Runtime error: {str(e)}\n\nFull trace:\n{traceback.format_exc()}"
-        print("=" * 80)
-        print(error_msg)
-        print("=" * 80)
-        torch.cuda.empty_cache()
-        return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-    except Exception as e:
-        error_msg = f"Generation failed: {str(e)}\n\nFull error:\n{traceback.format_exc()}"
-        print("=" * 80)
-        print("ERROR:")
-        print(error_msg)
-        print("=" * 80)
-        torch.cuda.empty_cache()
-        return gr.update(), gr.update(visible=False), gr.update(visible=True, value=error_msg)
-
-
-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):
-    """Generation with VRAM optimization"""
-    global last_fused, last_lora
-
-    try:
-        print("Loaded state dict:", loaded_state_dict)
-        print("Last LoRA:", last_lora, "| Current LoRA:", repo_name)
-
-        # Enhanced depth map generation
-        depth_image = enhanced_depth_map(face_image, face_detected)
-        
-        # CRITICAL FIX: Pipeline has 2 controlnets, must always pass 2 control images
-        if face_detected:
-            control_images = [face_kps, depth_image]
-            control_scales = [face_strength, depth_control_scale]
-        else:
-            # When no face detected, pass dummy black image for face controlnet with scale 0.0
-            dummy_face = Image.new('RGB', face_image.size, color=(0, 0, 0))
-            control_images = [dummy_face, depth_image]
-            control_scales = [0.0, depth_control_scale]  # Face control disabled
-
-        # Handle custom LoRA from HuggingFace
-        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 Exception("No .safetensors file found in this Hugging Face repository.")
-
-            weight_file = safetensors_files[0]
-            full_path_lora = hf_hub_download(repo_id=repo_id, filename=weight_file, repo_type="model")
-        else:
-            full_path_lora = loaded_state_dict
-
-        # LoRA loading
-        if last_lora != repo_name:
-            if last_fused:
-                pipe.unfuse_lora()
-                pipe.unload_lora_weights()
-                pipe.unload_textual_inversion()
-
-            try:
-                pipe.load_lora_weights(full_path_lora)
-                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
-                    )
-            except Exception as e:
-                raise Exception(f"Failed to load LoRA: {str(e)}")
-
-        print("✓ Processing embeddings...")
-        conditioning, pooled = compel(prompt)
-        negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
-
-        # VRAM OPTIMIZED: Reduced to 35 steps
-        num_inference_steps = 35
-        
-        print(f"✓ Generating image ({num_inference_steps} steps)...")
-        print(f"  Image size: {face_image.width}x{face_image.height}")
-        print(f"  Face detected: {face_detected}")
-        print(f"  Control images: {len(control_images)}")
-        print(f"  GPU Memory before generation: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
-        
-        try:
-            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=num_inference_steps,
-                guidance_scale=guidance_scale,
-                controlnet_conditioning_scale=control_scales,
-            ).images[0]
-        except torch.cuda.OutOfMemoryError as e:
-            print(f"  GPU Memory at error: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
-            raise Exception(f"CUDA out of memory during generation. Image size {face_image.width}x{face_image.height} is too large. Error: {str(e)}")
-        except RuntimeError as e:
-            if "out of memory" in str(e).lower():
-                print(f"  GPU Memory at error: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
-                raise Exception(f"GPU out of memory. Try smaller image or upgrade GPU. Error: {str(e)}")
-            else:
-                raise Exception(f"Runtime error during generation: {str(e)}")
-        except Exception as e:
-            raise Exception(f"Pipeline generation failed: {str(e)}")
-
-        # Post-processing detail enhancement
-        print("✓ Enhancing details...")
-        image = enhance_details(image, strength=1.15)
-
-        print(f"✓ Generation complete! GPU Memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
-        
-        last_lora = repo_name
-        return image
-        
-    except Exception as e:
-        raise Exception(f"Image generation failed: {str(e)}\n{traceback.format_exc()}")
-
-
-def detect_faces(face_image, use_multiple_faces=False):
-    """
-    Enhanced face detection with better filtering
-    """
-    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 []
-        
-        # Enhanced: Stricter quality filtering
-        filtered_faces = []
-        for face_info in face_info_list:
-            # Higher confidence threshold (0.6 instead of 0.5)
-            if 'det_score' in face_info and face_info['det_score'] > 0.6:
-                # Check minimum face size (80x80 instead of default)
-                bbox = face_info['bbox']
-                width = bbox[2] - bbox[0]
-                height = bbox[3] - bbox[1]
-                
-                if width >= 80 and height >= 80:
-                    # Check reasonable aspect ratio
-                    aspect_ratio = width / height
-                    if 0.6 <= aspect_ratio <= 1.4:
-                        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 (strict filtering)")
-            return []
-        
-        # Sort 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 resize_image_aspect_ratio(img, max_dim=768):
-    """
-    VRAM OPTIMIZED: Reduced max dimension to 768 to prevent CUDA OOM errors
-    """
-    width, height = img.size
-    aspect_ratio = width / height
-
-    if aspect_ratio >= 1:
-        new_width = min(max_dim, width)
-        new_height = int(new_width / aspect_ratio)
-    else:
-        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)
-
-
 # Build Gradio interface
 with gr.Blocks(css="custom.css") as demo:
     gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
     title = gr.HTML(
         """<h1><img src="https://i.imgur.com/DVoGw04.png">
-<span>LucasArts Style - VRAM Optimized âš¡<br><small style="
+<span>Face to All - Enhanced<br><small style="
     font-size: 13px;
     display: block;
     font-weight: normal;
     opacity: 0.75;
-">🔥 Fixed: Version compatibility + CUDA OOM errors resolved<br>
-✨ 768px max output | 35 inference steps | Detailed error messages<br>
-AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
+">🔥 Supports: No faces (landscape), Multiple faces, Improved quality, Custom LoRAs<br> diffusers InstantID + ControlNet</small></span></h1>""",
         elem_id="title",
     )
     selected_state = gr.State()
@@ -928,7 +632,7 @@ AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
                 photo = gr.Image(label="Upload a picture (with or without faces)", interactive=True, type="pil", height=300)
                 selected_loras = gr.Gallery(label="Selected LoRAs", height=80, show_share_button=False, visible=False, elem_id="gallery_selected")
                 gallery = gr.Gallery(
-                    label="LucasArts Style",
+                    label="Pick a style from the gallery",
                     allow_preview=False,
                     columns=4,
                     elem_id="gallery",
@@ -949,55 +653,24 @@ AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
                 interactive=False, label="Generated Image", elem_id="result-image", position=0.1
             )
             
-            error_message = gr.Textbox(
-                label="Error Details",
-                visible=False,
-                elem_id="error-message",
-                lines=5,
-                max_lines=10
-            )
-            
             with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
                 community_icon = gr.HTML(community_icon_html)
                 loading_icon = gr.HTML(loading_icon_html)
                 share_button = gr.Button("Share to community", elem_id="share-btn")
             
             with gr.Accordion("Advanced options", open=False):
-                gr.Markdown("""
-                ### VRAM Optimizations Active âš¡
-                - 🎯 768x768 face detection
-                - 📐 768px max output resolution (reduced for stability)
-                - âš¡ 35 inference steps (balanced quality/speed)
-                - 🔍 Detailed error messages with solutions
-                - 💾 Reduced memory usage
-                - ✅ Fixed version compatibility issues
-                """)
-                use_multiple_faces = gr.Checkbox(
-                    label="Process multiple faces separately", 
-                    value=False,
-                    info="Generate separate outputs for each detected face"
-                )
+                use_multiple_faces = gr.Checkbox(label="Use multiple faces (if detected)", value=False)
                 negative = gr.Textbox(label="Negative Prompt")
-                weight = gr.Slider(0, 10, value=1.0, step=0.1, label="LoRA weight")
-                face_strength = gr.Slider(
-                    0, 2, value=1.0, step=0.01, label="Face identity strength", 
-                    info="Higher = stronger face preservation (auto-adjusted)"
-                )
-                image_strength = gr.Slider(
-                    0, 1, value=0.15, step=0.01, label="Image structure strength", 
-                    info="Lower = more transformation"
-                )
-                guidance_scale = gr.Slider(
-                    0, 50, value=7.5, step=0.1, label="Guidance Scale",
-                    info="Auto-optimized per mode (7.5 faces, 8.5 landscapes)"
-                )
-                depth_control_scale = gr.Slider(
-                    0, 1, value=0.8, step=0.01, label="Depth ControlNet strength",
-                    info="3D structure preservation (auto-optimized)"
-                )
+                weight = gr.Slider(0, 10, value=0.9, step=0.1, label="LoRA weight")
+                face_strength = gr.Slider(0, 2, value=0.9, step=0.01, label="Face strength", 
+                                         info="Higher values increase face likeness (auto-adjusted for no-face images)")
+                image_strength = gr.Slider(0, 1, value=0.20, step=0.01, label="Image strength", 
+                                          info="Higher values increase similarity with original structure/colors")
+                guidance_scale = gr.Slider(0, 50, value=8, step=0.1, label="Guidance Scale")
+                depth_control_scale = gr.Slider(0, 1, value=0.8, step=0.01, label="Zoe Depth ControlNet strength")
             
             prompt_title = gr.Markdown(
-                value="### Click 'Run' to generate (VRAM optimized) âš¡",
+                value="### Click on a LoRA in the gallery to select it",
                 visible=True,
                 elem_id="selected_lora",
             )
@@ -1029,7 +702,7 @@ AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
         fn=run_lora,
         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, error_message],
+        outputs=[result, share_group],
     )
     
     button.click(
@@ -1040,7 +713,7 @@ AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
         fn=run_lora,
         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, error_message],
+        outputs=[result, share_group],
     )
     
     share_button.click(None, [], [], js=share_js)