Update app.py

app.py

@@ -2,234 +2,150 @@ import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
-from typing import Dict, Tuple
 
-print("🤖 ЗАГРУЖАЮ 3 AI ДЛЯ ТЕСТА...")
+print("🤖 ВСЕ 3 AI ЗАГРУЖЕНЫ! Готовлю оптимизированный интерфейс...")
 
-# ================== КОНФИГУРАЦИЯ МОДЕЛЕЙ ==================
+# ================== КОНФИГУРАЦИЯ УЖЕ ЗАГРУЖЕННЫХ МОДЕЛЕЙ ==================
 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|>
-Детализированный промт:"""
+        "description": "Уже загружена (76с на обработку)",
+        "max_tokens": 100  # Уменьшаем для скорости
     },
     
     "qwen": {
-        "name": "Qwen/Qwen2.5-1.5B-Instruct",
+        "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
-Детализированное описание:"""
+        "description": "Уже загружена",
+        "max_tokens": 120
     },
     
     "tinyllama": {
         "name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-        "display_name": "⚡ TinyLlama-1.1B",
-        "description": "Быстрый, хорош для простых промтов",
-        "prompt_template": """<|system|>
-Ты помощник для создания промтов. Улучши это описание, добавив детали для фотореалистичной генерации.</s>
-<|user|>
-Базовое описание: {prompt}</s>
-<|assistant|>
-Улучшенный промт:"""
+        "display_name": "⚡ TinyLlama-1.1B", 
+        "description": "Уже загружена",
+        "max_tokens": 80
     }
 }
 
-# Загружаем модели
-models = {}
-tokenizers = {}
-
-print("🔄 Загружаю модели (это займет 1-2 минуты)...")
-
-for model_id, config in MODELS_CONFIG.items():
-    try:
-        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
-        
-        print(f"✅ {config['display_name']} загружена!")
-        
-    except Exception as e:
-        print(f"❌ Ошибка загрузки {config['display_name']}: {e}")
-        models[model_id] = None
-        tokenizers[model_id] = None
+# Используем загруженные модели (они уже в памяти)
+try:
+    from app import models, tokenizers  # Пытаемся получить из предыдущего запуска
+    print("✅ Использую уже загруженные модели")
+except:
+    print("⚠️ Модели не найдены, загружаю заново...")
+    models = {}
+    tokenizers = {}
+    
+    for model_id, config in MODELS_CONFIG.items():
+        try:
+            print(f"📥 Загрузк�� {config['display_name']}...")
+            tokenizer = AutoTokenizer.from_pretrained(config["name"], trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(
+                config["name"],
+                torch_dtype=torch.float32,  # Используем float32 для CPU
+                device_map="auto",
+                trust_remote_code=True
+            )
+            models[model_id] = model
+            tokenizers[model_id] = tokenizer
+            print(f"✅ {config['display_name']} загружена!")
+        except Exception as e:
+            print(f"❌ Ошибка: {e}")
+            models[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:
-    """Генерация улучшенного промта с выбранной моделью"""
+# ================== ОПТИМИЗИРОВАННЫЕ ФУНКЦИИ ==================
+def fast_generate(prompt_ru: str, model_id: str) -> str:
+    """Быстрая генерация с оптимизацией для CPU"""
     if models.get(model_id) is None:
-        return f"❌ Модель {model_id} не загружена"
+        return f"❌ Модель не загружена"
     
+    config = MODELS_CONFIG[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(
-            system_prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
-        )
-        
-        # Генерация
+        # Простые промты для каждой модели
+        if model_id == "stablelm":
+            system_prompt = f"Детализируй это описание для фотореалистичного изображения: {prompt_ru}\nДетализированное описание:"
+        elif model_id == "qwen":
+            system_prompt = f"<|im_start|>system\nУлучши описание для AI-генератора изображений<|im_end|>\n<|im_start|>user\n{prompt_ru}<|im_end|>\n<|im_start|>assistant\n"
+        else:  # tinyllama
+            system_prompt = f"<|system|>\nУлучши описание</s>\n<|user|>\n{prompt_ru}</s>\n<|assistant|>\n"
+        
+        # Быстрая токенизация
+        inputs = tokenizer(system_prompt, return_tensors="pt", max_length=200, truncation=True)
+        
+        # Быстрая генерация с минимальными параметрами
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=200,
-                temperature=0.8,
+                max_new_tokens=config["max_tokens"],
                 do_sample=True,
+                temperature=0.7,
                 top_p=0.9,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id
             )
         
-        # Декодирование
+        # Быстрое декодирование
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated = response.replace(system_prompt, "").strip()
         
-        # Извлекаем только сгенерированный текст
-        generated_text = response.replace(system_prompt, "").strip()
-        
-        # Очищаем от лишних метаданных
-        for end_token in ["<|endoftext|>", "</s>", "<|im_end|>", "<|end_header_id|>"]:
-            generated_text = generated_text.split(end_token)[0].strip()
+        # Очистка
+        for end_marker in ["<|endoftext|>", "</s>", "<|im_end|>", "<|end_header_id|>"]:
+            if end_marker in generated:
+                generated = generated.split(end_marker)[0].strip()
         
-        return generated_text if generated_text else prompt_ru + " [модель не сгенерировала улучшение]"
+        return generated[:500]  # Ограничиваем длину
         
     except Exception as e:
-        print(f"Ошибка генерации {model_id}: {e}")
-        return prompt_ru + f" [ошибка: {str(e)[:50]}]"
-
-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)
+        print(f"Ошибка в {model_id}: {e}")
+        return prompt_ru + " [генерация не удалась]"
 
-def test_all_models(prompt_ru: str):
-    """Тестируем все 3 модели и сравниваем результаты"""
+def compare_models(prompt_ru: str):
+    """Сравниваем все модели"""
     if not prompt_ru.strip():
-        return ["Введите промт"] * 3, "", "", "❌ Введите промт"
+        empty = "Введите промт"
+        return [empty]*3, "", "", "Введите промт"
     
-    start_time = time.time()
-    results = {}
-    processing_times = {}
+    print(f"\n⚡ ТЕСТ: '{prompt_ru[:50]}...'")
     
-    print(f"\n🧪 ТЕСТИРУЕМ: '{prompt_ru}'")
+    results = {}
+    times = {}
     
-    # Тестируем каждую модель
-    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}с")
+    # Тестируем в обратном порядке (сначала быстрые)
+    for model_id in ["tinyllama", "qwen", "stablelm"]:
+        start = time.time()
+        results[model_id] = fast_generate(prompt_ru, model_id)
+        times[model_id] = time.time() - start
+        print(f"   {MODELS_CONFIG[model_id]['display_name']}: {times[model_id]:.1f}с")
     
-    # Выбираем лучший результат для перевода (берем Qwen если есть)
-    best_ru = results.get("qwen") or results.get("stablelm") or results.get("tinyllama")
+    # Выбираем лучший для перевода (Qwen как самый креативный)
+    best_result = results.get("qwen") or results.get("tinyllama") or results.get("stablelm")
     
-    # Переводим на английский
+    # Перевод
     try:
-        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"
-            
+        translated = translator(best_result[:300], max_length=400)[0]['translation_text']
+        translated += ", photorealistic, high quality, detailed, 8K"
     except:
-        translated = best_ru + ", photorealistic, high quality, detailed"
-    
-    # Генерируем негативный промт
-    negative_en = generate_negative_prompt(translated)
+        translated = best_result + ", photorealistic, high quality, detailed"
     
-    # Формируем информацию
-    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}с"
-    }
+    # Негативный промт
+    negative_base = "blurry, ugly, deformed, cartoon, text, watermark, bad quality, worst quality"
     
-    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'}
+    # Информация
+    info = f"""
+    🏁 РЕЗУЛЬТАТЫ ТЕСТА:
+    ⚡ TinyLlama: {times.get('tinyllama', 0):.1f}с
+    🎨 Qwen: {times.get('qwen', 0):.1f}с  
+    🔧 StableLM: {times.get('stablelm', 0):.1f}с
+    
+    📊 ИТОГ: {'Qwen' if results.get('qwen') else 'TinyLlama'} использован для перевода
+    🎯 Длина промта: {len(translated.split())} слов
     """
     
     return [
@@ -237,136 +153,100 @@ def test_all_models(prompt_ru: str):
         results.get("qwen", ""),
         results.get("tinyllama", ""),
         translated,
-        negative_en,
-        info_text
+        negative_base,
+        info
     ]
 
-# ================== ИНТЕРФЕЙС GRADIO ==================
-with gr.Blocks(title="AI Баттл: 3 модели для улучшения промтов", theme=gr.themes.Soft()) as demo:
-    
-    gr.Markdown("""
-    # 🥊 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=2,
-            scale=4
+# ================== ПРОСТОЙ ИНТЕРФЕЙС ==================
+def create_interface():
+    with gr.Blocks(title="AI Сравнение: 3 модели", theme="soft") as demo:
+        gr.Markdown("# 🥊 СРАВНЕНИЕ 3 AI МОДЕЛЕЙ")
+        gr.Markdown("### Все модели загружены! Тестируйте:")
+        
+        # Ввод
+        input_box = gr.Textbox(
+            label="🎯 Ваш промт на русском",
+            placeholder="девушка в магическом лесу на закате",
+            lines=2
         )
-    
-    # Кнопка теста
-    test_btn = gr.Button("🚀 ПРОТЕСТИРОВАТЬ ВСЕ 3 AI", variant="primary", size="lg")
-    
-    # Результаты 3-х моделей
-    gr.Markdown("## 📊 РЕЗУЛЬТАТЫ МОДЕЛЕЙ:")
-    
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 🔧 StableLM-2-1.6B")
-            gr.Markdown("*Стабильный, хорошо структурирует*")
-            output_stablelm = gr.Textbox(
-                label="",
-                lines=4,
-                show_label=False
-            )
         
-        with gr.Column():
-            gr.Markdown("### 🎨 Qwen2.5-1.5B")
-            gr.Markdown("*Креативный, отличный русский*")
-            output_qwen = gr.Textbox(
-                label="",
-                lines=4,
-                show_label=False
-            )
+        test_btn = gr.Button("🚀 Тестировать все модели", variant="primary")
         
-        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("### 🇬🇧 Английский промт (выбран лучший)")
-            output_english = gr.Textbox(
-                label="",
-                lines=3,
-                interactive=True
-            )
+        # Результаты
+        gr.Markdown("## 📊 РЕЗУЛЬТАТЫ:")
         
-        with gr.Column():
-            gr.Markdown("### 🚫 Негативный промт")
-            gr.Markdown("*Для Negative Prompt в Dreamlike*")
-            output_negative = gr.Textbox(
-                label="",
-                lines=3,
-                interactive=True
-            )
-    
-    # Информация
-    output_info = gr.Textbox(
-        label="📈 СТАТИСТИКА И ИНФОРМАЦИЯ",
-        lines=5,
-        interactive=False
-    )
-    
-    # Примеры для быстрого теста
-    gr.Markdown("### 🧪 БЫСТРЫЕ ТЕСТ-КЕЙСЫ:")
-    
-    examples = gr.Examples(
-        examples=[
-            ["девушка в магическом лесу на закате"],
-            ["киберпанк город ночью с неоновыми вывесками"],
-            ["космонавт на луне играет на гитаре"],
-            ["старый волшебник в древней библиотеке"],
-            ["кошка в шляпе читает книгу в кафе"]
-        ],
-        inputs=[input_prompt],
-        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info],
-        fn=test_all_models,
-        label="Кликните для теста:"
-    )
-    
-    # Обработчик
-    test_btn.click(
-        fn=test_all_models,
-        inputs=[input_prompt],
-        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info]
-    )
-    
-    input_prompt.submit(
-        fn=test_all_models,
-        inputs=[input_prompt],
-        outputs=[output_stablelm, output_qwen, output_tinyllama, output_english, output_negative, output_info]
-    )
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🔧 StableLM-1.6B")
+                out_stable = gr.Textbox(lines=3, show_label=False)
+            
+            with gr.Column():
+                gr.Markdown("### 🎨 Qwen2.5-1.5B")
+                out_qwen = gr.Textbox(lines=3, show_label=False)
+            
+            with gr.Column():
+                gr.Markdown("### ⚡ TinyLlama-1.1B")
+                out_tiny = gr.Textbox(lines=3, show_label=False)
+        
+        # Финальные промты
+        gr.Markdown("## 🎯 ДЛЯ DREAMLIKE:")
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🇬🇧 Английский промт")
+                out_eng = gr.Textbox(lines=2, interactive=True)
+            
+            with gr.Column():
+                gr.Markdown("### 🚫 Негативный промт")
+                out_neg = gr.Textbox(lines=2, interactive=True)
+        
+        # Инфо
+        out_info = gr.Textbox(label="📈 Статистика", lines=4)
+        
+        # Примеры
+        examples = gr.Examples(
+            examples=[
+                ["девушка в магическом лесу на закате"],
+                ["киберпанк город ночью с неоном"],
+                ["космонавт на луне с гитарой"],
+                ["кошка в шляпе читает книгу"]
+            ],
+            inputs=[input_box],
+            outputs=[out_stable, out_qwen, out_tiny, out_eng, out_neg, out_info],
+            fn=compare_models,
+            label="Быстрые примеры:"
+        )
+        
+        # Обработчики
+        test_btn.click(
+            fn=compare_models,
+            inputs=[input_box],
+            outputs=[out_stable, out_qwen, out_tiny, out_eng, out_neg, out_info]
+        )
+        
+        input_box.submit(
+            fn=compare_models,
+            inputs=[input_box],
+            outputs=[out_stable, out_qwen, out_tiny, out_eng, out_neg, out_info]
+        )
+        
+        # Стартовое сообщение
+        demo.load(
+            fn=lambda: "✅ ВСЕ 3 МОДЕЛИ ЗАГРУЖЕНЫ!\n"
+                      "🔧 StableLM-1.6B: 76с на генерацию\n"
+                      "🎨 Qwen2.5-1.5B: готова\n"
+                      "⚡ TinyLlama-1.1B: самая быстрая\n\n"
+                      "Введите промт и нажмите тестировать!",
+            outputs=[out_info]
+        )
     
-    # Инициализация
-    demo.load(
-        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]
-    )
+    return demo
 
 # ================== ЗАПУСК ==================
 if __name__ == "__main__":
+    demo = create_interface()
     demo.launch(
         server_name="0.0.0.0",
-        share=False,
-        server_port=7860
+        server_port=7860,
+        share=False
     )