Update app.py

app.py

@@ -1,322 +1,323 @@
 import gradio as gr
 import torch
 import time
-import os
-from typing import Dict, List
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-print("🚀 ЗАГРУЖАЮ СИСТЕМУ...")
+print("🚀 ЗАГРУЖАЮ ОПТИМИЗИРОВАННУЮ СИСТЕМУ...")
 
-# ================== ГЛОБАЛЬНЫЕ ПЕРЕМЕННЫЕ ДЛЯ КЭШИРОВАНИЯ ==================
-GLOBAL_MODELS = {}
-GLOBAL_TOKENIZERS = {}
+# ================== ИСПОЛЬЗУЕМ ТОЛЬКО TINYLLAMA ==================
+# Она быстрая, уже работает, занимает мало памяти
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# ================== ЗАГРУЗКА МОДЕЛЕЙ ОДИН РАЗ ==================
-def load_models_once():
-    """Загружает модели только один раз при старте"""
-    if GLOBAL_MODELS:
-        print("✅ Модели уже загружены (используем кэш)")
-        return True
+# Глобальные переменные для кэша
+_tokenizer = None
+_model = None
+_translator = None
+
+def load_system():
+    """Загружаем только одну модель - TinyLlama"""
+    global _tokenizer, _model, _translator
     
-    print("📦 Загружаю модели...")
+    print("📦 Загружаю TinyLlama (самая быстрая для CPU)...")
     
     try:
-        from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+        # Токенизатор
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        _tokenizer.pad_token = _tokenizer.eos_token
         
-        # Конфигурация моделей
-        models_to_load = {
-            "stablelm": ("stabilityai/stablelm-2-1_6b", "🔧 StableLM-2-1.6B"),
-            "qwen": ("Qwen/Qwen2.5-1.5B-Instruct", "🎨 Qwen2.5-1.5B"),
-            "tinyllama": ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "⚡ TinyLlama-1.1B")
-        }
+        # Модель с оптимизациями для CPU
+        _model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float32,  # Используем float32 для CPU
+            device_map="auto",
+            low_cpu_mem_usage=True,
+            offload_folder="./offload"  # Выгружаем на диск если не хватает RAM
+        )
         
-        # Загружаем каждую модель
-        for key, (model_name, display_name) in models_to_load.items():
-            try:
-                print(f"   Загрузка {display_name}...")
-                
-                # Токенизатор
-                tokenizer = AutoTokenizer.from_pretrained(
-                    model_name,
-                    trust_remote_code=True
-                )
-                
-                # Модель (используем float32 для CPU)
-                model = AutoModelForCausalLM.from_pretrained(
-                    model_name,
-                    torch_dtype=torch.float32,
-                    device_map="auto",
-                    trust_remote_code=True,
-                    low_cpu_mem_usage=True
-                )
-                
-                # Кэшируем
-                GLOBAL_TOKENIZERS[key] = tokenizer
-                GLOBAL_MODELS[key] = model
-                
-                print(f"   ✅ {display_name} загружена")
-                
-            except Exception as e:
-                print(f"   ❌ Ошибка загрузки {display_name}: {str(e)[:100]}")
-                GLOBAL_MODELS[key] = None
-                GLOBAL_TOKENIZERS[key] = None
+        print("✅ TinyLlama загружена и оптимизирована для CPU!")
         
-        # Загружаем переводчик
+        # Переводчик
         print("🌐 Загружаю переводчик...")
-        GLOBAL_MODELS["translator"] = pipeline(
+        _translator = pipeline(
             "translation_ru_to_en", 
-            model="Helsinki-NLP/opus-mt-ru-en"
+            model="Helsinki-NLP/opus-mt-ru-en",
+            device="cpu"
         )
-        print("✅ Переводчик загружен")
+        print("✅ Переводчик готов!")
         
         return True
         
     except Exception as e:
-        print(f"❌ Критическая ошибка загрузки: {e}")
+        print(f"❌ Ошибка загрузки: {e}")
         return False
 
-# ================== ФУНКЦИИ ГЕНЕРАЦИИ ==================
-def generate_with_model_fast(prompt: str, model_key: str) -> str:
-    """Быстрая генерация с кэшированной моделью"""
-    if model_key not in GLOBAL_MODELS or GLOBAL_MODELS[model_key] is None:
-        return f"[Модель {model_key} не загружена]"
+# ================== ПРАВИЛЬНЫЕ ПРОМТ-ШАБЛОНЫ ==================
+def get_enhancement_prompt(user_input: str) -> str:
+    """Создаем ПРАВИЛЬНЫЙ промт для улучшения, а не для рассказа"""
+    return f"""<|system|>
+Ты помощник для создания промтов к AI-генератору изображений Dreamlike Photoreal 2.0.
+ТВОЯ ЗАДАЧА: Взять короткое описание и добавить КОНКРЕТНЫЕ детали для фотореалистичной генерации.
+
+ДОБАВЬ:
+1. Описание внешности (если есть персонаж)
+2. Детали окружения
+3. Освещение и время суток
+4. Атмосферу и настроение
+5. Стиль изображения
+
+НЕ пиши рассказ, НЕ добавляй сюжет. Только описание для генератора изображений.
+
+Примеры правильных улучшений:
+• "девушка в лесу" → "Молодая девушка с рыжими волосами в белом платье стоит в магическом лесу на закате, золотые лучи солнца пробиваются сквозь листву, атмосферное освещение, фотореалистично, детализированно"
+• "город ночью" → "Футуристический город ночью с неоновыми вывесками, мокрые улицы отражают свет, киберпанк стиль, детализированная архитектура, фотореалистично"</s>
+<|user|>
+Улучши это описание для генерации фотореалистичного изображения: "{user_input}"</s>
+<|assistant|>
+Улучшенное описание:"""
+
+# ================== БЫСТРАЯ ГЕНЕРАЦИЯ ==================
+def enhance_prompt_fast(user_input: str) -> str:
+    """Быстро улучшаем промт с TinyLlama"""
+    if _model is None or _tokenizer is None:
+        return user_input + " [система не загружена]"
     
     try:
-        model = GLOBAL_MODELS[model_key]
-        tokenizer = GLOBAL_TOKENIZERS[model_key]
+        # Создаем правильный промт
+        system_prompt = get_enhancement_prompt(user_input)
         
-        # Простые промты для каждой модели
-        if model_key == "stablelm":
-            input_text = f"Детализируй: {prompt}\nДетали:"
-        elif model_key == "qwen":
-            input_text = f"<|im_start|>system\nДетализируй описание<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
-        else:  # tinyllama
-            input_text = f"<|system|>\nДетализируй<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
-        
-        # Генерация
-        inputs = tokenizer(input_text, return_tensors="pt", max_length=150, truncation=True)
+        # Токенизация (ограничиваем длину)
+        inputs = _tokenizer(
+            system_prompt,
+            return_tensors="pt",
+            max_length=256,
+            truncation=True,
+            padding=True
+        )
         
+        # Быстрая генерация с оптимизированными параметрами
         with torch.no_grad():
-            outputs = model.generate(
+            outputs = _model.generate(
                 **inputs,
-                max_new_tokens=80,
+                max_new_tokens=120,  # Не больше 120 новых токенов
+                temperature=0.8,
                 do_sample=True,
-                temperature=0.7,
-                pad_token_id=tokenizer.eos_token_id
+                top_p=0.9,
+                repetition_penalty=1.1,
+                pad_token_id=_tokenizer.eos_token_id,
+                eos_token_id=_tokenizer.eos_token_id,
+                no_repeat_ngram_size=3  # Избегаем повторений
             )
         
         # Декодируем
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        result = response.replace(input_text, "").strip()
+        response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Извлекаем только улучшенную часть
+        if "Улучшенное описание:" in response:
+            enhanced = response.split("Улучшенное описание:")[-1].strip()
+        else:
+            # Fallback: удаляем системный промт
+            enhanced = response.replace(system_prompt, "").strip()
         
-        # Очищаем
-        for marker in ["<|endoftext|>", "</s>", "<|im_end|>"]:
-            result = result.split(marker)[0].strip()
+        # Очищаем от метаданных
+        for marker in ["<|endoftext|>", "</s>", "<|assistant|>", "<|user|>", "<|system|>"]:
+            enhanced = enhanced.split(marker)[0].strip()
         
-        return result if result else prompt
+        # Если результат слишком короткий, добавляем базовые улучшения
+        if len(enhanced.split()) < 5:
+            enhanced = f"{user_input}, фотореалистично, высокое качество, детализированно, красивое освещение"
+        
+        return enhanced[:400]  # Ограничиваем длину
         
     except Exception as e:
-        return f"{prompt} [ошибка: {str(e)[:30]}]"
+        print(f"Ошибка генерации: {e}")
+        return f"{user_input}, фотореалистично, высокое качество, детализированно"
 
-def process_all_models(user_input: str):
-    """Обрабатывает промт всеми моделями"""
-    if not user_input.strip():
-        empty = "Введите промт"
-        return [empty]*3, "", "", "⏳ Введите промт..."
+def generate_negative_prompt_smart(english_prompt: str) -> str:
+    """Умный негативный промт"""
+    base = [
+        "blurry", "ugly", "deformed", "bad anatomy",
+        "bad proportions", "extra limbs", "mutated hands",
+        "poorly drawn hands", "text", "watermark", "signature",
+        "worst quality", "low quality", "jpeg artifacts"
+    ]
     
-    start_time = time.time()
+    prompt_lower = english_prompt.lower()
     
-    print(f"\n🧪 Обработка: '{user_input[:40]}...'")
+    # Добавляем контекстные исключения
+    if any(word in prompt_lower for word in ["realistic", "photoreal", "photo"]):
+        base.extend(["cartoon", "anime", "painting", "drawing", "illustration", "3d render"])
     
-    # Генерируем всеми моделями
-    results = {}
+    if any(word in prompt_lower for word in ["bright", "light", "sun"]):
+        base.extend(["dark", "dull", "underexposed", "gloomy"])
     
-    # TinyLlama (самая быстрая)
-    t1 = time.time()
-    results["tinyllama"] = generate_with_model_fast(user_input, "tinyllama")
-    time_tiny = time.time() - t1
+    if any(word in prompt_lower for word in ["portrait", "face", "person"]):
+        base.extend(["asymmetric eyes", "cloned face", "malformed limbs"])
     
-    # Qwen
-    t2 = time.time()
-    results["qwen"] = generate_with_model_fast(user_input, "qwen")
-    time_qwen = time.time() - t2
+    # Убираем дубликаты
+    unique = []
+    for term in base:
+        if term not in unique:
+            unique.append(term)
     
-    # StableLM (самая медленная)
-    t3 = time.time()
-    results["stablelm"] = generate_with_model_fast(user_input, "stablelm")
-    time_stable = time.time() - t3
+    return ", ".join(unique)
+
+def process_pipeline(user_input: str):
+    """Полный пайплайн обработки"""
+    start_time = time.time()
     
-    # Выбираем лучший результат для перевода (Qwen приоритет)
-    best_ru = results.get("qwen") or results.get("tinyllama") or results.get("stablelm")
+    if not user_input.strip():
+        return "Введите промт", "", "", "⏳ Введите промт"
     
-    # Переводим
+    print(f"⚡ Обработка: {user_input}")
+    
+    # 1. Улучшаем промт (максимум 5 секунд)
+    gen_start = time.time()
+    enhanced_ru = enhance_prompt_fast(user_input)
+    gen_time = time.time() - gen_start
+    
+    # 2. Переводим
+    trans_start = time.time()
     try:
-        translator = GLOBAL_MODELS.get("translator")
-        if translator:
-            translated = translator(best_ru[:250])[0]['translation_text']
+        if _translator:
+            translated = _translator(enhanced_ru[:300])[0]['translation_text']
         else:
-            translated = best_ru
+            translated = enhanced_ru
     except:
-        translated = best_ru
+        translated = enhanced_ru
+    trans_time = time.time() - trans_start
     
-    # Добавляем ключевые слова
-    translated += ", photorealistic, high quality, detailed, 8K"
+    # 3. Улучшаем английский для Dreamlike
+    dreamlike_keywords = ", photorealistic, high quality, detailed, 8K, professional photography"
+    if "photorealistic" not in translated.lower():
+        translated += dreamlike_keywords
     
-    # Негативный промт
-    negative = "blurry, ugly, deformed, cartoon, text, watermark, bad quality, worst quality"
+    # 4. Негативный промт
+    negative_start = time.time()
+    negative_en = generate_negative_prompt_smart(translated)
+    negative_time = time.time() - negative_start
     
-    # Информация
+    # 5. Статистика
     total_time = time.time() - start_time
+    
     info = f"""
-    ⏱️ ВРЕМЯ ОБРАБОТКИ:
-    • ⚡ TinyLlama: {time_tiny:.1f} сек
-    • 🎨 Qwen: {time_qwen:.1f} сек
-    • 🔧 StableLM: {time_stable:.1f} сек
+    ⚡ БЫСТРАЯ ОБРАБОТКА:
+    • Генерация: {gen_time:.1f} сек
+    • Перевод: {trans_time:.1f} сек
+    • Негативный промт: {negative_time:.1f} сек
     • Всего: {total_time:.1f} сек
     
     📊 СТАТИСТИКА:
-    • Использован: {'Qwen' if results.get('qwen') else 'TinyLlama'}
-    • Длина английского: {len(translated.split())} слов
-    • Статус: ✅ Готово
+    • Модель: TinyLlama-1.1B (оптимизирована для CPU)
+    • Длина промта: {len(translated.split())} слов
+    • Статус: ✅ Готово к использованию в Dreamlike
     """
     
-    return [
-        results.get("stablelm", ""),
-        results.get("qwen", ""),
-        results.get("tinyllama", ""),
-        translated,
-        negative,
-        info.strip()
-    ]
+    return enhanced_ru, translated, negative_en, info.strip()
 
-# ================== ИНТЕРФЕЙС GRADIO ==================
-def create_simple_interface():
-    with gr.Blocks(title="AI Баттл: 3 Модели") as demo:
+# ================== ПРОСТОЙ ИНТЕРФЕЙС ==================
+with gr.Blocks(title="AI Улучшитель промтов для Dreamlike", theme=gr.themes.Soft()) as demo:
+    
+    # Заголовок
+    gr.Markdown("""
+    # 🚀 УЛЬТРА-БЫСТРЫЙ AI УЛУЧШИТЕЛЬ ПРОМТОВ
+    ### Оптимизирован для CPU • TinyLlama • Мгновенная обработка
+    """)
+    
+    # Статус загрузки
+    status = gr.Textbox(
+        label="🔄 Статус системы",
+        value="Загружаю систему...",
+        interactive=False
+    )
+    
+    # Основной интерфейс
+    with gr.Column(visible=False) as main_ui:
+        # Ввод
+        input_prompt = gr.Textbox(
+            label="🎯 Введите промт на русском",
+            placeholder="Пример: девушка в магическом лесу на закате",
+            lines=2
+        )
         
-        # Заголовок
-        gr.Markdown("""
-        # 🥊 БАТТЛ AI МОДЕЛЕЙ
-        ### StableLM-1.6B vs Qwen2.5-1.5B vs TinyLlama-1.1B
-        """)
+        # Кнопка
+        process_btn = gr.Button("⚡ УЛУЧШИТЬ ПРОМТ", variant="primary", size="lg")
         
-        # Статус загрузки
-        load_status = gr.Textbox(
-            label="🔄 Статус системы",
-            value="Загружаю модели... (это займет ~2 минуты)",
-            interactive=False
+        # Результаты
+        gr.Markdown("### 📝 УЛУЧШЕННЫЙ ПРОМТ (русский)")
+        output_ru = gr.Textbox(
+            label="",
+            lines=3,
+            interactive=True
         )
         
-        # Основной интерфейс появится после загрузки
-        def build_interface_after_load():
-            # Проверяем загрузку
-            success = load_models_once()
-            
-            if not success:
-                return "❌ Ошибка загрузки моделей", gr.update(visible=False), gr.update(visible=False)
-            
-            # Считаем загруженные модели
-            loaded_count = sum(1 for m in ["stablelm", "qwen", "tinyllama"] if GLOBAL_MODELS.get(m) is not None)
+        gr.Markdown("### 🌐 ДЛЯ DREAMLIKE (английский)")
+        with gr.Row():
+            output_en = gr.Textbox(
+                label="🎨 Основной промт",
+                lines=3,
+                interactive=True
+            )
             
-            return f"✅ Загружено моделей: {loaded_count}/3 + Переводчик", gr.update(visible=True), gr.update(visible=True)
+            output_neg = gr.Textbox(
+                label="🚫 Negative Prompt",
+                lines=3,
+                interactive=True
+            )
         
-        # Создаем скрытые элементы
-        main_interface = gr.Column(visible=False)
-        examples_section = gr.Column(visible=False)
+        # Информация
+        output_info = gr.Textbox(
+            label="📊 СТАТИСТИКА",
+            lines=4,
+            interactive=False
+        )
         
-        # После загрузки показываем основной интерфейс
-        demo.load(
-            fn=build_interface_after_load,
-            outputs=[load_status, main_interface, examples_section]
+        # Примеры
+        gr.Markdown("### 🧪 БЫСТРЫЕ ПРИМЕРЫ")
+        examples = gr.Examples(
+            examples=[
+                ["девушка в магическом лесу на закате"],
+                ["киберпанк улица ночью с неоновыми вывесками"],
+                ["космонавт играет на гитаре на луне"],
+                ["кошка в шляпе читает газету в кафе"],
+                ["дракон на вершине горы"]
+            ],
+            inputs=[input_prompt],
+            outputs=[output_ru, output_en, output_neg, output_info],
+            fn=process_pipeline,
+            label="Кликните для теста:"
         )
         
-        # ОСНОВНОЙ ИНТЕРФЕЙС (показывается после загрузки)
-        with main_interface:
-            # Ввод
-            input_prompt = gr.Textbox(
-                label="🎯 Введите промт на русском",
-                placeholder="Пример: девушка в магическом лесу на закате",
-                lines=2
-            )
-            
-            # Кнопка
-            test_btn = gr.Button("🚀 Тестировать все 3 модели", variant="primary", size="lg")
-            
-            # Результаты
-            gr.Markdown("### 📊 РЕЗУЛЬТАТЫ МОДЕЛЕЙ")
-            
-            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_english = gr.Textbox(lines=2, interactive=True)
-                
-                with gr.Column():
-                    gr.Markdown("**🚫 Негативный промт**")
-                    out_negative = gr.Textbox(lines=2, interactive=True)
-            
-            # Информация
-            out_info = gr.Textbox(
-                label="📈 СТАТИСТИКА И ИНФОРМАЦИЯ",
-                lines=5,
-                interactive=False
-            )
-            
-            # Обработчики
-            def process_input(prompt):
-                return process_all_models(prompt)
-            
-            test_btn.click(
-                fn=process_input,
-                inputs=[input_prompt],
-                outputs=[out_stable, out_qwen, out_tiny, out_english, out_negative, out_info]
-            )
-            
-            input_prompt.submit(
-                fn=process_input,
-                inputs=[input_prompt],
-                outputs=[out_stable, out_qwen, out_tiny, out_english, out_negative, out_info]
-            )
+        # Обработчики
+        def process_and_update(prompt):
+            return process_pipeline(prompt)
         
-        # ПРИМЕРЫ (показываются после загрузки)
-        with examples_section:
-            gr.Markdown("### 🧪 БЫСТРЫЕ ПРИМЕРЫ ДЛЯ ТЕСТА")
-            
-            examples = gr.Examples(
-                examples=[
-                    ["девушка в магическом лесу на закате"],
-                    ["киберпанк город ночью с неоновыми вывесками"],
-                    ["космонавт на луне играет на гитаре"],
-                    ["старый волшебник в древней библиотеке"],
-                    ["кошка в шляпе читает книгу в кафе"]
-                ],
-                inputs=[input_prompt],
-                outputs=[out_stable, out_qwen, out_tiny, out_english, out_negative, out_info],
-                fn=process_input,
-                label="Кликните на пример:"
-            )
+        process_btn.click(
+            fn=process_and_update,
+            inputs=[input_prompt],
+            outputs=[output_ru, output_en, output_neg, output_info]
+        )
+        
+        input_prompt.submit(
+            fn=process_and_update,
+            inputs=[input_prompt],
+            outputs=[output_ru, output_en, output_neg, output_info]
+        )
     
-    return demo
+    # Функция инициализации
+    def initialize_system():
+        if load_system():
+            return "✅ СИСТЕМА ГОТОВА! TinyLlama загружена и оптимизирована.\n⚡ Введите промт и нажмите кнопку.", gr.update(visible=True)
+        else:
+            return "❌ Ошибка загрузки. П��пробуйте перезапустить Space.", gr.update(visible=False)
+    
+    # Загружаем при старте
+    demo.load(
+        fn=initialize_system,
+        outputs=[status, main_ui]
+    )
 
 # ================== ЗАПУСК ==================
 if __name__ == "__main__":
-    print("🚀 Создаю интерфейс...")
-    demo = create_simple_interface()
-    
-    # Убираем лишние параметры из launch
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860