Update app.py

app.py

@@ -2,267 +2,371 @@ import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
-import re
+from typing import Dict, Tuple
 
-# Загружаем модель для улучшения промтов
-print("🔄 Загружаю AI для улучшения промтов...")
+print("🤖 ЗАГРУЖАЮ 3 AI ДЛЯ ТЕСТА...")
 
-# Используем маленькую, но умную модель
-MODEL_NAME = "IlyaGusev/rubert_telegram_headlines"  # Хорошо понимает русский
-# ИЛИ альтернатива: "cointegrated/rubert-tiny2" - еще меньше
+# ================== КОНФИГУРАЦИЯ МОДЕЛЕЙ ==================
+MODELS_CONFIG = {
+    "stablelm": {
+        "name": "stabilityai/stablelm-2-1_6b",
+        "display_name": "🔧 StableLM-2-1.6B",
+        "description": "Стабильный, хорошо структурирует",
+        "prompt_template": """<|start_header_id|>system<|end_header_id|>
+Ты эксперт по созданию детализированных промтов для фотореалистичной генерации.
+Возьми исходную идею и добавь КОНКРЕТНЫЕ детали: описание объектов, освещение, атмосферу, мелкие детали.
+Сделай описание максимально наглядным для AI-генератора.<|eot_id|>
+<|start_header_id|>user<|end_header_id|>
+Исходная идея: {prompt}<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>
+Детализированный промт:"""
+    },
+    
+    "qwen": {
+        "name": "Qwen/Qwen2.5-1.5B-Instruct",
+        "display_name": "🎨 Qwen2.5-1.5B",
+        "description": "Креативный, отличный русский",
+        "prompt_template": """<|im_start|>system
+Ты художник-концептуалист. Придумай КРАСИВОЕ и ДЕТАЛЬНОЕ описание для генерации фотореалистичного изображения.
+Добавь: детали объектов, освещение, атмосферу, настроение, мелкие элементы фона.
+Будь максимально креативным!<|im_end|>
+<|im_start|>user
+Идея: {prompt}<|im_end|>
+<|im_start|>assistant
+Детализированное описание:"""
+    },
+    
+    "tinyllama": {
+        "name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        "display_name": "⚡ TinyLlama-1.1B",
+        "description": "Быстрый, хорош для простых промтов",
+        "prompt_template": """<|system|>
+Ты помощник для создания промтов. Улучши это описание, добавив детали для фотореалистичной генерации.</s>
+<|user|>
+Базовое описание: {prompt}</s>
+<|assistant|>
+Улучшенный промт:"""
+    }
+}
 
-try:
-    # Загружаем модель для генерации текста
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-    print(f"✅ Модель {MODEL_NAME} загружена!")
-except Exception as e:
-    print(f"⚠️ Не удалось загрузить основную модель: {e}")
-    print("🔄 Использую альтернативную модель...")
-    MODEL_NAME = "cointegrated/rubert-tiny2"
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-    print(f"✅ Альтернативная модель загружена!")
+# Загружаем модели
+models = {}
+tokenizers = {}
 
-# Загружаем переводчик
-print("🔄 Загружаю переводчик...")
-translator = pipeline("translation_ru_to_en", model="Helsinki-NLP/opus-mt-ru-en", device="cpu")
-print("✅ Переводчик готов!")
+print("🔄 Загружаю модели (это займет 1-2 минуты)...")
 
-def ai_enhance_prompt(prompt_ru):
-    """AI улучшает промт, генерируя новые детали"""
+for model_id, config in MODELS_CONFIG.items():
     try:
-        # Создаем промт для AI
-        system_prompt = """Ты эксперт по созданию промтов для генерации изображений. 
-        Улучши этот промт, добавив детали: опиши освещение, время суток, атмосферу, стиль, детали окружения.
+        print(f"📥 Загрузка {config['display_name']}...")
+        
+        # Загружаем токенизатор
+        tokenizer = AutoTokenizer.from_pretrained(
+            config["name"],
+            trust_remote_code=True
+        )
+        
+        # Загружаем модель
+        model = AutoModelForCausalLM.from_pretrained(
+            config["name"],
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        # Настраиваем специальные токены
+        if model_id == "tinyllama":
+            tokenizer.pad_token = tokenizer.eos_token
         
-        Исходный промт: """
+        models[model_id] = model
+        tokenizers[model_id] = tokenizer
         
-        full_prompt = system_prompt + prompt_ru + "\nУлучшенный промт: "
+        print(f"✅ {config['display_name']} загружена!")
+        
+    except Exception as e:
+        print(f"❌ Ошибка загрузки {config['display_name']}: {e}")
+        models[model_id] = None
+        tokenizers[model_id] = None
+
+# Загружаем переводчик
+print("🌐 Загружаю переводчик...")
+translator = pipeline("translation_ru_to_en", model="Helsinki-NLP/opus-mt-ru-en")
+print("✅ Переводчик готов!")
+
+# ================== ФУНКЦИИ ГЕНЕРАЦИИ ==================
+def generate_with_model(prompt_ru: str, model_id: str) -> str:
+    """Генерация улучшенного промта с выбранной моделью"""
+    if models.get(model_id) is None:
+        return f"❌ Модель {model_id} не загружена"
+    
+    model = models[model_id]
+    tokenizer = tokenizers[model_id]
+    config = MODELS_CONFIG[model_id]
+    
+    try:
+        # Формируем промт
+        system_prompt = config["prompt_template"].format(prompt=prompt_ru)
         
-        # Токенизируем
-        inputs = tokenizer(full_prompt, return_tensors="pt", max_length=512, truncation=True)
+        # Токенизация
+        inputs = tokenizer(
+            system_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        )
         
-        # Генерируем улучшенный промт
+        # Генерация
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=100,
+                max_new_tokens=200,
                 temperature=0.8,
                 do_sample=True,
                 top_p=0.9,
                 repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
         
-        # Декодируем ответ
+        # Декодирование
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Извлекаем улучшенный промт (удаляем системный промт)
-        if "Улучшенный промт:" in response:
-            enhanced = response.split("Улучшенный промт:")[1].strip()
-        else:
-            # Fallback: добавляем базовые улучшения
-            enhanced = prompt_ru + ", фотореалистично, высокое качество, детализированно, красивое освещение"
+        # Извлекаем только сгенерированный текст
+        generated_text = response.replace(system_prompt, "").strip()
         
-        # Убедимся, что промт не слишком короткий
-        if len(enhanced.split()) < 8:
-            enhanced += ", профессиональная фотография, чёткие детали, приятная атмосфера"
+        # Очищаем от лишних метаданных
+        for end_token in ["<|endoftext|>", "</s>", "<|im_end|>", "<|end_header_id|>"]:
+            generated_text = generated_text.split(end_token)[0].strip()
         
-        return enhanced
+        return generated_text if generated_text else prompt_ru + " [модель не сгенерировала улучшение]"
         
     except Exception as e:
-        print(f"Ошибка в AI улучшении: {e}")
-        # Fallback на ручное улучшение
-        return prompt_ru + ", фотореалистично, высокое качество, 8K, профессиональная фотография"
+        print(f"Ошибка генерации {model_id}: {e}")
+        return prompt_ru + f" [ошибка: {str(e)[:50]}]"
 
-def generate_negative_prompt(english_prompt):
-    """AI генерирует негативный промт на основе английского промта"""
-    try:
-        # Анализируем промт и генерируем противоположные термины
-        prompt_lower = english_prompt.lower()
-        
-        negative_terms = []
-        
-        # Базовые термины для исключения
-        base_negative = ["blurry", "ugly", "deformed", "bad quality", "worst quality"]
-        negative_terms.extend(base_negative)
-        
-        # Определяем, что исключать на основе промта
-        if "realistic" in prompt_lower or "photoreal" in prompt_lower:
-            negative_terms.extend(["cartoon", "anime", "painting", "drawing", "illustration"])
-        
-        if "bright" in prompt_lower or "sunny" in prompt_lower:
-            negative_terms.extend(["dark", "dull", "underexposed"])
-        
-        if "portrait" in prompt_lower or "face" in prompt_lower:
-            negative_terms.extend(["deformed face", "asymmetric eyes", "bad anatomy", "extra limbs"])
-        
-        if "detailed" in prompt_lower:
-            negative_terms.extend(["simple", "plain", "boring"])
-        
-        if "beautiful" in prompt_lower or "gorgeous" in prompt_lower:
-            negative_terms.extend(["unattractive", "messy", "dirty"])
-        
-        # Добавляем артефакты, которые всегда исключаем
-        negative_terms.extend(["text", "watermark", "signature", "logo", "username"])
-        
-        # Убираем дубликаты
-        unique_terms = []
-        for term in negative_terms:
-            if term not in unique_terms:
-                unique_terms.append(term)
-        
-        return ", ".join(unique_terms)
-        
-    except Exception as e:
-        print(f"Ошибка в генерации негативного промта: {e}")
-        return "blurry, ugly, deformed, cartoon, text, watermark, bad quality"
+def generate_negative_prompt(english_prompt: str) -> str:
+    """Генерация негативного промта"""
+    base_negative = [
+        "blurry", "ugly", "deformed", "bad quality", "worst quality",
+        "cartoon", "anime", "painting", "drawing", "illustration",
+        "text", "watermark", "signature", "logo", "username",
+        "bad anatomy", "extra limbs", "mutated hands"
+    ]
+    
+    # Анализируем промт
+    prompt_lower = english_prompt.lower()
+    
+    if any(word in prompt_lower for word in ["realistic", "photoreal", "photo"]):
+        base_negative.extend(["unrealistic", "artificial", "CGI"])
+    
+    if any(word in prompt_lower for word in ["bright", "sunny", "light"]):
+        base_negative.extend(["dark", "dull", "underexposed"])
+    
+    if any(word in prompt_lower for word in ["detailed", "intricate"]):
+        base_negative.extend(["simple", "plain", "boring"])
+    
+    # Убираем дубликаты
+    unique_terms = list(dict.fromkeys(base_negative))
+    return ", ".join(unique_terms)
 
-def process_full_pipeline(user_prompt):
-    """Полный пайплайн обработки промта"""
+def test_all_models(prompt_ru: str):
+    """Тестируем все 3 модели и сравниваем результаты"""
+    if not prompt_ru.strip():
+        return ["Введите промт"] * 3, "", "", "❌ Введите промт"
+    
     start_time = time.time()
+    results = {}
+    processing_times = {}
     
-    if not user_prompt.strip():
-        return "Введите промт", "", "", "0.0s"
+    print(f"\n🧪 ТЕСТИРУЕМ: '{prompt_ru}'")
     
-    print(f"🚀 Обрабатываю промт: {user_prompt}")
+    # Тестируем каждую модель
+    for model_id in ["stablelm", "qwen", "tinyllama"]:
+        model_start = time.time()
+        
+        if models.get(model_id) is None:
+            results[model_id] = f"❌ {MODELS_CONFIG[model_id]['display_name']} не загружена"
+        else:
+            enhanced = generate_with_model(prompt_ru, model_id)
+            results[model_id] = enhanced
+        
+        processing_times[model_id] = time.time() - model_start
+        print(f"   {MODELS_CONFIG[model_id]['display_name']}: {processing_times[model_id]:.1f}с")
     
-    # 1. AI улучшает русский промт
-    print("🔄 Этап 1: AI улучшает промт...")
-    enhanced_ru = ai_enhance_prompt(user_prompt)
+    # Выбираем лучший результат для перевода (берем Qwen если есть)
+    best_ru = results.get("qwen") or results.get("stablelm") or results.get("tinyllama")
     
-    # 2. Переводим на английский
-    print("🔄 Этап 2: Перевод на английский...")
+    # Переводим на английский
     try:
-        translated_result = translator(enhanced_ru, max_length=400)[0]
-        optimized_en = translated_result['translation_text']
+        translated = translator(best_ru, max_length=400)[0]['translation_text']
+        
+        # Добавляем ключевые слова для Dreamlike
+        if "photorealistic" not in translated.lower():
+            translated += ", photorealistic"
+        if "high quality" not in translated.lower():
+            translated += ", high quality"
+        if "detailed" not in translated.lower():
+            translated += ", detailed"
+            
     except:
-        # Fallback перевод
-        optimized_en = enhanced_ru + ", photorealistic, high quality, detailed"
-    
-    # 3. AI генерирует негативный промт
-    print("🔄 Этап 3: Генерация негативного промта...")
-    negative_en = generate_negative_prompt(optimized_en)
-    
-    # 4. Улучшаем английский промт для Dreamlike
-    if "photorealistic" not in optimized_en.lower():
-        optimized_en += ", photorealistic"
-    if "high quality" not in optimized_en.lower():
-        optimized_en += ", high quality"
-    if "detailed" not in optimized_en.lower():
-        optimized_en += ", detailed"
+        translated = best_ru + ", photorealistic, high quality, detailed"
     
-    # 5. Добавляем стиль для Dreamlike Photoreal
-    if "dreamlike" not in optimized_en.lower():
-        optimized_en += ", dreamlike aesthetic"
+    # Генерируем негативный промт
+    negative_en = generate_negative_prompt(translated)
     
-    processing_time = time.time() - start_time
+    # Формируем информацию
+    total_time = time.time() - start_time
+    time_info = {
+        "stablelm": f"{processing_times.get('stablelm', 0):.1f}с",
+        "qwen": f"{processing_times.get('qwen', 0):.1f}с",
+        "tinyllama": f"{processing_times.get('tinyllama', 0):.1f}с",
+        "total": f"{total_time:.1f}с"
+    }
     
-    print(f"✅ Готово! Время: {processing_time:.1f}с")
+    info_text = f"""
+    ⏱️ ВРЕМЯ ОБРАБОТКИ:
+    • StableLM-1.6B: {time_info['stablelm']}
+    • Qwen2.5-1.5B: {time_info['qwen']} 
+    • TinyLlama-1.1B: {time_info['tinyllama']}
+    • Всего: {time_info['total']}
+    ✅ Использован для перевода: {MODELS_CONFIG.get('qwen', {}).get('display_name', 'Qwen') if results.get('qwen') else 'StableLM'}
+    """
     
-    return (
-        enhanced_ru,
-        optimized_en,
+    return [
+        results.get("stablelm", ""),
+        results.get("qwen", ""),
+        results.get("tinyllama", ""),
+        translated,
         negative_en,
-        f"⏱️ {processing_time:.1f} сек | ✅ AI улучшила промт"
-    )
+        info_text
+    ]
 
-# Создаем простой и понятный интерфейс
-with gr.Blocks(title="AI Улучшитель промтов для Dreamlike", theme=gr.themes.Soft()) as demo:
+# ================== ИНТЕРФЕЙС GRADIO ==================
+with gr.Blocks(title="AI Баттл: 3 модели для улучшения промтов", theme=gr.themes.Soft()) as demo:
+    
     gr.Markdown("""
-    # 🎨 AI Улучшитель промтов для Dreamlike Photoreal 2.0
-    ### Нейросеть сама улучшит ваш промт, переведет и создаст негативный промт
+    # 🥊 AI БАТТЛ: 3 МОДЕЛИ ДЛЯ УЛУЧШЕНИЯ ПРОМТОВ
+    ### Сравниваем StableLM-1.6B vs Qwen2.5-1.5B vs TinyLlama-1.1B
     """)
     
+    # Статус загрузки моделей
+    loaded_models = [m for m, val in models.items() if val is not None]
+    gr.Markdown(f"### ✅ ЗАГРУЖЕНО МОДЕЛЕЙ: {len(loaded_models)}/3")
+    
     # Ввод
     with gr.Row():
         input_prompt = gr.Textbox(
-            label="📝 Введите ваш промт на русском",
-            placeholder="Пример: девушка в лесу на закате",
-            lines=3,
+            label="🎯 Введите промт для теста",
+            placeholder="Пример: девушка в магическом лесу на закате",
+            lines=2,
             scale=4
         )
     
-    # Кнопка
-    submit_btn = gr.Button("🚀 Улучшить промт с помощью AI", variant="primary", size="lg")
+    # Кнопка теста
+    test_btn = gr.Button("🚀 ПРОТЕСТИРОВАТЬ ВСЕ 3 AI", variant="primary", size="lg")
+    
+    # Результаты 3-х моделей
+    gr.Markdown("## 📊 РЕЗУЛЬТАТЫ МОДЕЛЕЙ:")
     
-    # Результаты
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 🇷🇺 Улучшенный промт (AI-генерация)")
-            output_ru = gr.Textbox(
+            gr.Markdown("### 🔧 StableLM-2-1.6B")
+            gr.Markdown("*Стабильный, хорошо структурирует*")
+            output_stablelm = gr.Textbox(
                 label="",
                 lines=4,
-                interactive=True
+                show_label=False
             )
+        
+        with gr.Column():
+            gr.Markdown("### 🎨 Qwen2.5-1.5B")
+            gr.Markdown("*Креативный, отличный русский*")
+            output_qwen = gr.Textbox(
+                label="",
+                lines=4,
+                show_label=False
+            )
+        
+        with gr.Column():
+            gr.Markdown("### ⚡ TinyLlama-1.1B")
+            gr.Markdown("*Быстрый, хорош для простых промтов*")
+            output_tinyllama = gr.Textbox(
+                label="",
+                lines=4,
+                show_label=False
+            )
+    
+    # Финальные промты для использования
+    gr.Markdown("## 🎯 ФИНАЛЬНЫЕ ПРОМТЫ ДЛЯ DREAMLIKE:")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 🇬🇧 Промт для Dreamlike (английский)")
-            output_en = gr.Textbox(
+            gr.Markdown("### 🇬🇧 Английский промт (выбран лучший)")
+            output_english = gr.Textbox(
                 label="",
-                lines=4,
+                lines=3,
+                interactive=True
+            )
+        
+        with gr.Column():
+            gr.Markdown("### 🚫 Негативный промт")
+            gr.Markdown("*Для Negative Prompt в Dreamlike*")
+            output_negative = gr.Textbox(
+                label="",
+                lines=3,
                 interactive=True
             )
-    
-    gr.Markdown("### 🚫 Негативный промт (английский)")
-    gr.Markdown("**Скопируйте это в Negative Prompt:**")
-    output_negative = gr.Textbox(
-        label="",
-        lines=3,
-        interactive=True
-    )
     
     # Информация
-    info_text = gr.Textbox(
-        label="📊 Информация",
-        value="🔄 Готов к работе! Введите промт и нажмите кнопку.",
+    output_info = gr.Textbox(
+        label="📈 СТАТИСТИКА И ИНФОРМАЦИЯ",
+        lines=5,
         interactive=False
     )
     
-    # Примеры
-    gr.Markdown("### 🧪 Примеры для теста:")
+    # Примеры для быстрого теста
+    gr.Markdown("### 🧪 БЫСТРЫЕ ТЕСТ-КЕЙСЫ:")
+    
     examples = gr.Examples(
         examples=[
-            ["девушка в лесу на закате"],
-            ["космонавт на луне с гитарой"],
-            ["киберпанк город ночью с неоном"],
-            ["старый волшебник в библиотеке"],
-            ["кошка в шляпе читает книгу"]
+            ["девушка в магическом лесу на закате"],
+            ["киберпанк город ночью с неоновыми вывесками"],
+            ["космонавт на луне играет на гитаре"],
+            ["старый волшебник в древней библиотеке"],
+            ["кошка в шляпе читает книгу в кафе"]
         ],
         inputs=[input_prompt],
-        label="Кликните на пример:"
+        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info],
+        fn=test_all_models,
+        label="Кликните для теста:"
     )
     
-    # Обработчики
-    def process_and_update(prompt):
-        enhanced, english, negative, info = process_full_pipeline(prompt)
-        return enhanced, english, negative, info
-    
-    submit_btn.click(
-        fn=process_and_update,
+    # Обработчик
+    test_btn.click(
+        fn=test_all_models,
         inputs=[input_prompt],
-        outputs=[output_ru, output_en, output_negative, info_text]
+        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info]
     )
     
     input_prompt.submit(
-        fn=process_and_update,
+        fn=test_all_models,
         inputs=[input_prompt],
-        outputs=[output_ru, output_en, output_negative, info_text]
+        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info]
     )
     
-    # Приветственное сообщение
+    # Инициализация
     demo.load(
-        fn=lambda: "✅ Система готова! AI модели загружены. Введите промт.",
-        outputs=[info_text]
+        fn=lambda: f"✅ Система готова! Загружено моделей: {len(loaded_models)}/3\n"
+                  f"• 🔧 StableLM-2-1.6B: {'✅' if models.get('stablelm') else '❌'}\n"
+                  f"• 🎨 Qwen2.5-1.5B: {'✅' if models.get('qwen') else '❌'}\n"
+                  f"• ⚡ TinyLlama-1.1B: {'✅' if models.get('tinyllama') else '❌'}",
+        outputs=[output_info]
     )
 
-# Запуск
+# ================== ЗАПУСК ==================
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         share=False,
-        server_port=7860,
-        show_error=True
+        server_port=7860
     )