Add support for local LoRA from /workspace/loras/

app.py

@@ -148,25 +148,113 @@ def resize_image(input_image, max_size=1024):
 # LORA FUNCTIONS
 # =================================================================
 
-# Список доступных LoRA
-AVAILABLE_LORAS = {
+# Папка для локальных LoRA
+LOCAL_LORA_DIR = "/workspace/loras"
+
+# Базовые LoRA из HuggingFace Hub
+HUB_LORAS = {
     "Realism": {
         "repo": "flymy-ai/qwen-image-realism-lora",
         "trigger": "Super Realism portrait of",
-        "weights": "pytorch_lora_weights.safetensors"
+        "weights": "pytorch_lora_weights.safetensors",
+        "source": "hub"
     },
     "Anime": {
         "repo": "alfredplpl/qwen-image-modern-anime-lora",
         "trigger": "Japanese modern anime style, ",
-        "weights": "pytorch_lora_weights.safetensors"
+        "weights": "pytorch_lora_weights.safetensors",
+        "source": "hub"
     },
     "Analog Film": {
         "repo": "janekm/analog_film",
         "trigger": "fifthel",
-        "weights": "converted_complete.safetensors"
+        "weights": "converted_complete.safetensors",
+        "source": "hub"
     }
 }
 
+def scan_local_loras():
+    """
+    Сканирует папку /workspace/loras на наличие .safetensors файлов
+    Возвращает dict с найденными LoRA
+    """
+    local_loras = {}
+    
+    if not os.path.exists(LOCAL_LORA_DIR):
+        logger.info(f"  Local LoRA directory not found: {LOCAL_LORA_DIR}")
+        return local_loras
+    
+    logger.info(f"  Scanning local LoRA directory: {LOCAL_LORA_DIR}")
+    
+    try:
+        for file in os.listdir(LOCAL_LORA_DIR):
+            if file.endswith('.safetensors'):
+                lora_name = os.path.splitext(file)[0]  # Имя без расширения
+                local_path = os.path.join(LOCAL_LORA_DIR, file)
+                
+                # Добавляем в список
+                local_loras[lora_name] = {
+                    "path": local_path,
+                    "trigger": "",  # Без trigger word для локальных
+                    "weights": file,
+                    "source": "local"
+                }
+                
+                logger.info(f"    ✓ Found local LoRA: {lora_name} ({file})")
+    
+    except Exception as e:
+        logger.warning(f"  Error scanning local LoRA directory: {e}")
+    
+    return local_loras
+
+# Сканируем локальные LoRA
+logger.info("\nScanning for LoRA models...")
+LOCAL_LORAS = scan_local_loras()
+
+# Объединяем Hub и локальные LoRA
+AVAILABLE_LORAS = {**HUB_LORAS, **LOCAL_LORAS}
+
+if LOCAL_LORAS:
+    logger.info(f"  ✓ Found {len(LOCAL_LORAS)} local LoRA(s)")
+logger.info(f"  Total available LoRAs: {len(AVAILABLE_LORAS)}")
+
+def load_lora_weights(pipeline, lora_name, lora_scale, hf_token):
+    """
+    Загружает LoRA веса в pipeline
+    Поддерживает как Hub LoRA так и локальные
+    """
+    if lora_name == "None" or lora_name not in AVAILABLE_LORAS:
+        return None
+    
+    lora_info = AVAILABLE_LORAS[lora_name]
+    
+    try:
+        if lora_info['source'] == 'hub':
+            # Загрузка с HuggingFace Hub
+            logger.info(f"  Loading LoRA from Hub: {lora_info['repo']}")
+            pipeline.load_lora_weights(
+                lora_info['repo'],
+                weight_name=lora_info.get('weights', 'pytorch_lora_weights.safetensors'),
+                token=hf_token
+            )
+        else:
+            # Загрузка локального файла
+            logger.info(f"  Loading local LoRA: {lora_info['path']}")
+            pipeline.load_lora_weights(
+                lora_info['path'],
+                adapter_name=lora_name
+            )
+        
+        # Устанавливаем scale
+        if hasattr(pipeline, 'set_adapters'):
+            pipeline.set_adapters([lora_name], adapter_weights=[lora_scale])
+        
+        return lora_info.get('trigger', '')
+        
+    except Exception as e:
+        logger.error(f"  ❌ Error loading LoRA {lora_name}: {e}")
+        return None
+
 # =================================================================
 # GENERATION FUNCTIONS
 # =================================================================
@@ -212,20 +300,14 @@ def generate_text2img(
     
     try:
         # Загружаем LoRA если выбрана
-        if lora_name != "None" and lora_name in AVAILABLE_LORAS:
-            lora_info = AVAILABLE_LORAS[lora_name]
-            logger.info(f"  Loading LoRA: {lora_info['repo']}")
-            
-            pipe_txt2img.load_lora_weights(
-                lora_info['repo'],
-                weight_name=lora_info.get('weights', 'pytorch_lora_weights.safetensors'),
-                token=hf_token
-            )
+        trigger_word = None
+        if lora_name != "None":
+            trigger_word = load_lora_weights(pipe_txt2img, lora_name, lora_scale, hf_token)
             
-            # Добавляем trigger word
-            if lora_info['trigger']:
-                prompt = lora_info['trigger'] + prompt
-                logger.info(f"  Added trigger: {lora_info['trigger']}")
+            # Добавляем trigger word если есть
+            if trigger_word:
+                prompt = trigger_word + prompt
+                logger.info(f"  Added trigger: {trigger_word}")
         
         generator = torch.Generator(device=device).manual_seed(seed)
         
@@ -291,15 +373,13 @@ def generate_img2img(
             raise gr.Error("Image2Image pipeline not available")
         
         # Загружаем LoRA если выбрана
-        if lora_name != "None" and lora_name in AVAILABLE_LORAS:
-            lora_info = AVAILABLE_LORAS[lora_name]
-            pipe_img2img.load_lora_weights(
-                lora_info['repo'],
-                weight_name=lora_info.get('weights', 'pytorch_lora_weights.safetensors'),
-                token=hf_token
-            )
-            if lora_info['trigger']:
-                prompt = lora_info['trigger'] + prompt
+        trigger_word = None
+        if lora_name != "None":
+            trigger_word = load_lora_weights(pipe_img2img, lora_name, lora_scale, hf_token)
+            
+            # Добавляем trigger word если есть
+            if trigger_word:
+                prompt = trigger_word + prompt
         
         generator = torch.Generator(device=device).manual_seed(seed)
         
@@ -341,7 +421,9 @@ css = """
 """
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
+    lora_choices = ["None"] + list(AVAILABLE_LORAS.keys())
+    
+    gr.Markdown(f"""
     # 🎨 Qwen Soloband - Image2Image + LoRA
     
     **Продвинутая модель генерации** с поддержкой Text-to-Image, Image-to-Image и LoRA стилей.
@@ -349,11 +431,13 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     ### ✨ Возможности:
     - 🖼️ **Text-to-Image** - Генерация из текста, разрешения до 2048×2048
     - 🔄 **Image-to-Image** - Модификация изображений с контролем strength (0.0-1.0)
-    - 🎭 **LoRA Support** - Динамическая загрузка стилей (Realism, Anime, Film)
+    - 🎭 **LoRA Support** - {len(AVAILABLE_LORAS)} доступных стилей (Hub + локальные)
     - 🔌 **Full API** - Все функции доступны через API
     - ⚡ **Optimized** - VAE tiling/slicing, правильный QwenImageImg2ImgPipeline
     
     **Модель**: [Gerchegg/Qwen-Soloband-Diffusers](https://huggingface.co/Gerchegg/Qwen-Soloband-Diffusers)
+    
+    💡 **Local LoRAs**: Положите .safetensors файлы в `/workspace/loras/` - они появятся автоматически!
     """)
     
     with gr.Tabs() as tabs:
@@ -387,8 +471,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         
                         t2i_lora = gr.Radio(
                             label="LoRA Style",
-                            choices=["None"] + list(AVAILABLE_LORAS.keys()),
-                            value="None"
+                            choices=lora_choices,
+                            value="None",
+                            info=f"Hub: {len(HUB_LORAS)}, Local: {len(LOCAL_LORAS)}"
                         )
                         t2i_lora_scale = gr.Slider(label="LoRA Strength", minimum=0.0, maximum=2.0, step=0.1, value=1.0)
                 
@@ -431,8 +516,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         
                         i2i_lora = gr.Radio(
                             label="LoRA Style",
-                            choices=["None"] + list(AVAILABLE_LORAS.keys()),
-                            value="None"
+                            choices=lora_choices,
+                            value="None",
+                            info=f"Hub: {len(HUB_LORAS)}, Local: {len(LOCAL_LORAS)}"
                         )
                         i2i_lora_scale = gr.Slider(label="LoRA Strength", minimum=0.0, maximum=2.0, step=0.1, value=1.0)