Update app.py

app.py

@@ -38,7 +38,7 @@ from compel import Compel, ReturnedEmbeddingsType
 
 from gradio_imageslider import ImageSlider
 
-# Load LoRA configurations
+# Load LoRA configurations - now only LucasArts style
 with open("sdxl_loras.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -123,8 +123,9 @@ et = time.time()
 print('Loading VAE took: ', et - st, 'seconds')
 
 st = time.time()
+# CHANGED: Using AlbedoBase XL v2.1 for better quality
 pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0", 
+    "frankjoshua/albedobaseXL_v21", 
     vae=vae,
     controlnet=[identitynet, zoedepthnet],
     torch_dtype=torch.float16
@@ -132,7 +133,8 @@ pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
 
 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)
+# IMPROVED: Higher IP adapter scale for better face preservation
+pipe.set_ip_adapter_scale(1.0)
 et = time.time()
 print('Loading pipeline took: ', et - st, 'seconds')
 
@@ -157,10 +159,10 @@ last_lora = ""
 last_fused = False
 lora_archive = "/data"
 
-# Improved face detection with multi-face support
+# Enhanced face detection with better face quality filtering
 def detect_faces(face_image, use_multiple_faces=False):
     """
-    Detect faces in the image
+    Detect faces in the image with quality filtering
     Returns: list of face info dictionaries, or empty list if no faces
     """
     try:
@@ -170,47 +172,56 @@ def detect_faces(face_image, use_multiple_faces=False):
             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)
-        face_info_list = sorted(
-            face_info_list, 
+        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(face_info_list)} faces")
-            return face_info_list
+            print(f"Detected {len(filtered_faces)} high-quality faces")
+            return filtered_faces
         else:
-            print(f"Using largest face (detected {len(face_info_list)} total)")
-            return [face_info_list[0]]
+            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(face_info_list):
+def process_face_embeddings_separately(face_info_list):
     """
-    Process face embeddings - average multiple faces or return single face
+    Process face embeddings separately for multi-face generation
+    Returns: list of individual face embeddings
     """
     if not face_info_list:
-        return None
-    
-    if len(face_info_list) == 1:
-        return face_info_list[0]['embedding']
+        return []
     
-    # Average embeddings for multiple faces
     embeddings = [face_info['embedding'] for face_info in face_info_list]
-    avg_embedding = np.mean(embeddings, axis=0)
-    return avg_embedding
+    return embeddings
 
 def create_face_kps_image(face_image, face_info_list):
     """
-    Create keypoints image from face info
+    Create keypoints image from face info with enhanced visibility
     """
     if not face_info_list:
         return face_image
     
-    # For multiple faces, draw all keypoints
+    # For multiple faces, draw all keypoints with different colors
     if len(face_info_list) > 1:
         return draw_multiple_kps(face_image, [f['kps'] for f in face_info_list])
     else:
@@ -218,7 +229,7 @@ def create_face_kps_image(face_image, face_info_list):
 
 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
+    Draw keypoints for multiple faces with enhanced visibility
     """
     stickwidth = 4
     limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
@@ -226,12 +237,14 @@ def draw_multiple_kps(image_pil, kps_list, color_list=[(255, 0, 0), (0, 255, 0),
     w, h = image_pil.size
     out_img = np.zeros([h, w, 3])
     
-    for kps in kps_list:
+    for idx, kps in enumerate(kps_list):
         kps = np.array(kps)
+        # Use different colors for different faces
+        color_offset = idx % len(color_list)
         
         for i in range(len(limbSeq)):
             index = limbSeq[i]
-            color = color_list[index[0]]
+            color = color_list[(index[0] + color_offset) % len(color_list)]
             
             x = kps[index][:, 0]
             y = kps[index][:, 1]
@@ -245,7 +258,7 @@ 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 = color_list[(idx_kp + color_offset) % len(color_list)]
             x, y = kp
             out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
     
@@ -260,9 +273,9 @@ def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, i
 
     for lora_list in lora_defaults:
         if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
-            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)
+            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)
             depth_control_scale = lora_list.get("depth_control_scale", 0.8)
             negative = lora_list.get("negative", "")
     
@@ -308,10 +321,7 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
              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
+    Enhanced run_lora with improved face preservation and landscape mode
     """
     print("Custom LoRA:", custom_lora)
     custom_lora_path = custom_lora[0] if custom_lora else None
@@ -320,26 +330,22 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     st = time.time()
     face_image = resize_image_aspect_ratio(face_image)
 
-    # Validate selection FIRST before any array access
-    if not selected_state and not custom_lora_path:
-        raise gr.Error("❌ You must select a style before generating")
-    
-    # Validate index is within bounds
-    if not custom_lora_path and (selected_state_index < 0 or selected_state_index >= len(sdxl_loras)):
-        raise gr.Error(f"❌ Invalid selection index: {selected_state_index}")
-
     # 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)
+        # CHANGED: 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)")
+        print(f"Processing with {len(face_info_list)} face(s) separately")
+        
+        # For multiple faces, we'll generate with the primary face (largest)
+        face_emb = face_embeddings[0]
     else:
         face_emb = None
         face_kps = face_image
-        print("No faces detected - using landscape/depth mode only")
+        print("No faces detected - using enhanced landscape/depth mode")
 
     et = time.time()
     print('Face processing took:', et - st, 'seconds')
@@ -356,18 +362,27 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
                 if prompt_full:
                     prompt = prompt_full.replace("<subject>", prompt)
 
+    # Add LucasArts trigger word if not present
+    if "lucasarts artstyle" not in prompt.lower():
+        prompt = f"{prompt}, lucasarts artstyle"
+
     print("Prompt:", prompt)
     if prompt == "":
-        prompt = "a beautiful scene" if not face_detected else "a person"
+        prompt = "a beautiful cinematic scene" if not face_detected else "a person in cinematic lighting"
     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
+        # Enhanced negative prompt
+        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:
+    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:
@@ -381,15 +396,22 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     et = time.time()
     print('Prompt processing took:', et - st, 'seconds')
 
-    # Adjust parameters based on face detection
+    # IMPROVED: Better parameter adjustment for face/landscape modes
     if not face_detected:
-        # For landscape/no face mode, reduce face strength and increase depth control
+        # Enhanced landscape mode parameters
         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")
+        depth_control_scale = 1.0  # Maximum depth control for landscapes
+        image_strength = 0.25  # Higher structure preservation
+        print("Adjusted parameters for enhanced landscape mode")
+    else:
+        # Enhanced face preservation
+        face_strength = max(face_strength, 1.0)  # Ensure strong face preservation
+        depth_control_scale = max(depth_control_scale, 0.8)  # Good depth control
+        print("Adjusted parameters for enhanced face preservation")
 
     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,
@@ -401,7 +423,7 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
 run_lora.zerogpu = True
 
 
-@spaces.GPU(duration=75)
+@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):
@@ -409,15 +431,17 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
 
     print("Loaded state dict:", loaded_state_dict)
     print("Last LoRA:", last_lora, "| Current LoRA:", repo_name)
-    print("LoRA scale:", lora_scale, "Type:", type(lora_scale))
 
-    # Prepare control images and scales based on face detection
+    # IMPROVED: Better control image preparation
+    depth_image = zoe(face_image)
+    
     if face_detected:
-        control_images = [face_kps, zoe(face_image)]
+        # Face mode: use both face keypoints and depth
+        control_images = [face_kps, depth_image]
         control_scales = [face_strength, depth_control_scale]
     else:
-        # Only use depth control for landscapes
-        control_images = [zoe(face_image)]
+        # Landscape mode: only depth control with enhanced parameters
+        control_images = [depth_image]
         control_scales = [depth_control_scale]
 
     # Handle custom LoRA from HuggingFace
@@ -438,46 +462,17 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     # Improved LoRA loading and caching
     if last_lora != repo_name:
         if last_fused:
-            try:
-                # Unload previous LoRA adapter
-                pipe.delete_adapters(["style_lora"])
-                print("Unloaded previous LoRA adapter")
-            except Exception as e:
-                print(f"Warning: Could not unload previous adapter: {e}")
-                # Try alternative method
-                try:
-                    pipe.unload_lora_weights()
-                    print("Unloaded LoRA weights via alternative method")
-                except:
-                    pass
-            
-            try:
-                pipe.unload_textual_inversion()
-                print("Unloaded textual inversion")
-            except Exception as e:
-                print(f"Warning: Could not unload textual inversion: {e}")
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            pipe.unload_textual_inversion()
 
         # Load LoRA with better error handling
         try:
-            # For diffusers >= 0.27, load_lora_weights expects different parameters
-            if full_path_lora.endswith('.safetensors') or full_path_lora.endswith('.bin'):
-                # Single file loading
-                import os
-                lora_dir = os.path.dirname(full_path_lora)
-                lora_file = os.path.basename(full_path_lora)
-                print(f"Loading LoRA from: {lora_dir}/{lora_file}")
-                pipe.load_lora_weights(lora_dir, weight_name=lora_file, adapter_name="style_lora")
-            else:
-                # Directory loading
-                print(f"Loading LoRA from directory: {full_path_lora}")
-                pipe.load_lora_weights(full_path_lora, adapter_name="style_lora")
-            
-            # Set adapter scale instead of fusing
-            pipe.set_adapters(["style_lora"], adapter_weights=[float(lora_scale)])
-            print(f"LoRA loaded and adapter set with scale: {lora_scale}")
+            pipe.load_lora_weights(full_path_lora)
+            pipe.fuse_lora(lora_scale)
             last_fused = True
             
-            # Handle pivotal tuning embeddings
+            # Handle pivotal tuning embeddings (if needed for future LoRAs)
             is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
             if is_pivotal:
                 text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
@@ -496,20 +491,15 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
                     tokenizer=pipe.tokenizer_2
                 )
         except Exception as e:
-            import traceback
-            full_error = traceback.format_exc()
             print(f"Error loading LoRA: {e}")
-            print(f"Full traceback:\n{full_error}")
-            print(f"Full path attempted: {full_path_lora}")
-            print(f"LoRA scale attempted: {lora_scale} (type: {type(lora_scale)})")
-            raise gr.Error(f"Failed to load LoRA: {str(e)}\n\nPath: {full_path_lora}\nScale: {lora_scale}")
+            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
+    # IMPROVED: Enhanced generation parameters for better quality
+    num_inference_steps = 50  # Increased for better quality
     
     print("Generating image...")
     image = pipe(
@@ -521,7 +511,7 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
         height=face_image.height,
         image_embeds=face_emb if face_detected else None,
         image=face_image,
-        strength=1-image_strength,
+        strength=1-image_strength,  # Higher strength = more transformation
         control_image=control_images,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
@@ -656,12 +646,12 @@ 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>Face to All - Enhanced<br><small style="
+<span>LucasArts Style - Enhanced Face Preservation<br><small style="
     font-size: 13px;
     display: block;
     font-weight: normal;
     opacity: 0.75;
-">🔥 Supports: No faces (landscape), Multiple faces, Improved quality, Custom LoRAs<br> diffusers InstantID + ControlNet</small></span></h1>""",
+">🔥 Improved: Better face identity preservation, Enhanced landscape mode, Multiple face support<br>AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
         elem_id="title",
     )
     selected_state = gr.State()
@@ -673,7 +663,7 @@ with gr.Blocks(css="custom.css") as demo:
                 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="Pick a style from the gallery",
+                    label="LucasArts Style",
                     allow_preview=False,
                     columns=4,
                     elem_id="gallery",
@@ -700,18 +690,32 @@ with gr.Blocks(css="custom.css") as demo:
                 share_button = gr.Button("Share to community", elem_id="share-btn")
             
             with gr.Accordion("Advanced options", open=False):
-                use_multiple_faces = gr.Checkbox(label="Use multiple faces (if detected)", value=False)
+                use_multiple_faces = gr.Checkbox(
+                    label="Process multiple faces separately", 
+                    value=False,
+                    info="Generate separate outputs for each detected face"
+                )
                 negative = gr.Textbox(label="Negative Prompt")
-                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")
+                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 for landscapes)"
+                )
+                image_strength = gr.Slider(
+                    0, 1, value=0.15, step=0.01, label="Image structure strength", 
+                    info="Lower = more transformation, Higher = more original structure"
+                )
+                guidance_scale = gr.Slider(
+                    0, 50, value=7.5, step=0.1, label="Guidance Scale",
+                    info="How closely to follow the prompt"
+                )
+                depth_control_scale = gr.Slider(
+                    0, 1, value=0.8, step=0.01, label="Depth ControlNet strength",
+                    info="3D structure preservation (auto-maximized for landscapes)"
+                )
             
             prompt_title = gr.Markdown(
-                value="### Click on a LoRA in the gallery to select it",
+                value="### Click 'Run' to generate with LucasArts style",
                 visible=True,
                 elem_id="selected_lora",
             )
@@ -761,4 +765,4 @@ with gr.Blocks(css="custom.css") as demo:
     demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras])
 
 demo.queue(default_concurrency_limit=None, api_open=True)
-demo.launch()
+demo.launch(share=True)