Update app.py

app.py

@@ -8,7 +8,7 @@ import os
 from PIL import Image
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, list_repo_files
 from safetensors.torch import load_file
 import requests
 import re
@@ -37,7 +37,50 @@ try:
     print(f"Successfully loaded {len(flux_loras_raw)} LoRAs from JSON")
 except Exception as e:
     print(f"Error loading flux_loras.json: {e}")
-    flux_loras_raw = []
+    print("Using sample LoRA data instead...")
+    # Sample LoRA data with working repositories
+    flux_loras_raw = [
+        {
+            "image": "https://huggingface.co/alvdansen/flux-koda/resolve/main/images/photo-1586902197503-e71026292412.jpeg",
+            "title": "Flux Koda",
+            "repo": "alvdansen/flux-koda",
+            "trigger_word": "flmft style",
+            "weights": "flux_lora.safetensors",
+            "likes": 100
+        },
+        {
+            "image": "https://huggingface.co/multimodalart/flux-tarot-v1/resolve/main/images/e5f2761e5a474e52ab11b1c9246c9a30.png",
+            "title": "Tarot Cards",
+            "repo": "multimodalart/flux-tarot-v1",
+            "trigger_word": "in the style of TOK a trtcrd tarot style",
+            "weights": "flux_tarot_v1_lora.safetensors",
+            "likes": 90
+        },
+        {
+            "image": "https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style/resolve/main/d13591878de740648a8f29b836e16ff2.jpeg",
+            "title": "Paper Cutout",
+            "repo": "Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style",
+            "trigger_word": "Paper Cutout Style",
+            "weights": "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors",
+            "likes": 80
+        },
+        {
+            "image": "https://huggingface.co/alvdansen/frosting_lane_flux/resolve/main/images/content%20-%202024-08-11T010011.238.jpeg",
+            "title": "Frosting Lane",
+            "repo": "alvdansen/frosting_lane_flux",
+            "trigger_word": "frstingln illustration",
+            "weights": "flux_lora_frosting_lane_flux_000002500.safetensors",
+            "likes": 70
+        },
+        {
+            "image": "https://huggingface.co/davisbro/flux-watercolor/resolve/main/images/wc2.png",
+            "title": "Watercolor",
+            "repo": "davisbro/flux-watercolor",
+            "trigger_word": "watercolor style",
+            "weights": "flux_watercolor.safetensors",
+            "likes": 60
+        }
+    ]
 # Global variables for LoRA management
 current_lora = None
 lora_cache = {}
@@ -45,10 +88,45 @@ lora_cache = {}
 def load_lora_weights(repo_id, weights_filename):
     """Load LoRA weights from HuggingFace"""
     try:
-        if repo_id not in lora_cache:
+        # First try with the specified filename
+        try:
             lora_path = hf_hub_download(repo_id=repo_id, filename=weights_filename)
-            lora_cache[repo_id] = lora_path
-        return lora_cache[repo_id]
+            if repo_id not in lora_cache:
+                lora_cache[repo_id] = lora_path
+            return lora_path
+        except Exception as e:
+            print(f"Failed to load {weights_filename}, trying to find alternative LoRA files...")
+            
+            # If the specified file doesn't exist, try to find any .safetensors file
+            from huggingface_hub import list_repo_files
+            try:
+                files = list_repo_files(repo_id)
+                safetensors_files = [f for f in files if f.endswith(('.safetensors', '.bin')) and 'lora' in f.lower()]
+                
+                if not safetensors_files:
+                    # Try without 'lora' in filename
+                    safetensors_files = [f for f in files if f.endswith('.safetensors')]
+                
+                if safetensors_files:
+                    # Try the first available file
+                    for file in safetensors_files:
+                        try:
+                            print(f"Trying alternative file: {file}")
+                            lora_path = hf_hub_download(repo_id=repo_id, filename=file)
+                            if repo_id not in lora_cache:
+                                lora_cache[repo_id] = lora_path
+                            print(f"Successfully loaded alternative LoRA file: {file}")
+                            return lora_path
+                        except:
+                            continue
+                            
+                print(f"No suitable LoRA files found in {repo_id}")
+                return None
+                
+            except Exception as list_error:
+                print(f"Error listing files in repo {repo_id}: {list_error}")
+                return None
+                
     except Exception as e:
         print(f"Error loading LoRA from {repo_id}: {e}")
         return None
@@ -74,21 +152,36 @@ def get_huggingface_lora(link):
             model_card = ModelCard.load(link)
             trigger_word = model_card.data.get("instance_prompt", "")
             
-            fs = HfFileSystem()
-            list_of_files = fs.ls(link, detail=False)
-            safetensors_file = None
+            # Try to find the correct safetensors file
+            files = list_repo_files(link)
+            safetensors_files = [f for f in files if f.endswith('.safetensors')]
             
-            for file in list_of_files:
-                if file.endswith(".safetensors") and "lora" in file.lower():
-                    safetensors_file = file.split("/")[-1]
-                    break
-            
-            if not safetensors_file:
-                safetensors_file = "pytorch_lora_weights.safetensors"
+            # Prioritize files with 'lora' in the name
+            lora_files = [f for f in safetensors_files if 'lora' in f.lower()]
+            if lora_files:
+                safetensors_file = lora_files[0]
+            elif safetensors_files:
+                safetensors_file = safetensors_files[0]
+            else:
+                # Try .bin files as fallback
+                bin_files = [f for f in files if f.endswith('.bin') and 'lora' in f.lower()]
+                if bin_files:
+                    safetensors_file = bin_files[0]
+                else:
+                    safetensors_file = "pytorch_lora_weights.safetensors"  # Default fallback
             
+            print(f"Found LoRA file: {safetensors_file} in {link}")
             return split_link[1], safetensors_file, trigger_word
+            
         except Exception as e:
-            raise Exception(f"Error loading LoRA: {e}")
+            print(f"Error in get_huggingface_lora: {e}")
+            # Try basic detection
+            try:
+                files = list_repo_files(link)
+                safetensors_file = next((f for f in files if f.endswith('.safetensors')), "pytorch_lora_weights.safetensors")
+                return split_link[1], safetensors_file, ""
+            except:
+                raise Exception(f"Error loading LoRA: {e}")
     else:
         raise Exception("Invalid HuggingFace repository format")
 
@@ -130,12 +223,27 @@ def classify_gallery(flux_loras):
     try:
         sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
         gallery_items = []
+        
         for item in sorted_gallery:
             if "image" in item and "title" in item:
-                gallery_items.append((item["image"], item["title"]))
+                image_url = item["image"]
+                title = item["title"]
+                
+                # If image is a local file path that might not exist, use a placeholder URL
+                if isinstance(image_url, str) and (image_url.startswith("/home/") or image_url.startswith("samples/") or not image_url.startswith("http")):
+                    print(f"Replacing local/invalid image path: {image_url}")
+                    # Use a more reliable placeholder
+                    image_url = f"https://via.placeholder.com/512x512/E0E7FF/818CF8?text={title.replace(' ', '+')}"
+                    
+                gallery_items.append((image_url, title))
+        
+        if not gallery_items:
+            print("No gallery items found after filtering")
+            return [], sorted_gallery
+        
         return gallery_items, sorted_gallery
     except Exception as e:
-        print(f"Error loading gallery: {e}")
+        print(f"Error in classify_gallery: {e}")
         return [], []
 
 def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
@@ -477,7 +585,7 @@ with gr.Blocks(css=css) as demo:
                 custom_model = gr.Textbox(
                     label="Or enter a custom HuggingFace FLUX LoRA", 
                     placeholder="e.g., username/lora-name",
-                    visible=False
+                    visible=True
                 )
                 custom_model_card = gr.HTML(visible=False)
                 custom_model_button = gr.Button("Remove custom LoRA", visible=False)