Update app.py

app.py

@@ -3,35 +3,41 @@ import torch
 import time
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-print("🚀 ЗАГРУЖАЮ ИСПРАВЛЕННУЮ СИСТЕМУ...")
+print("🚀 ЗАГРУЖАЮ ОПТИМИЗИРОВАННУЮ СИСТЕМУ...")
 
-# ================== КОНФИГУРАЦИЯ ==================
+# ================== ИСПОЛЬЗУЕМ ТОЛЬКО TINYLLAMA ==================
+# Она быстрая, уже работает, занимает мало памяти
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# Глобальные переменные
+# Глобальные переменные для кэша
 _tokenizer = None
 _model = None
 _translator = None
 
 def load_system():
-    """Загружаем систему"""
+    """Загружаем только одну модель - TinyLlama"""
     global _tokenizer, _model, _translator
     
-    print("📦 Загружаю TinyLlama...")
+    print("📦 Загружаю TinyLlama (самая быстрая для CPU)...")
     
     try:
+        # Токенизатор
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         _tokenizer.pad_token = _tokenizer.eos_token
         
+        # Модель с оптимизациями для CPU
         _model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float32,  # Используем float32 для CPU
             device_map="auto",
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            offload_folder="./offload"  # Выгружаем на диск если не хватает RAM
         )
         
-        print("✅ TinyLlama загружена!")
+        print("✅ TinyLlama загружена и оптимизирована для CPU!")
         
+        # Переводчик
+        print("🌐 Загружаю переводчик...")
         _translator = pipeline(
             "translation_ru_to_en", 
             model="Helsinki-NLP/opus-mt-ru-en",
@@ -42,222 +48,184 @@ def load_system():
         return True
         
     except Exception as e:
-        print(f"❌ Ошибка: {e}")
+        print(f"❌ Ошибка загрузки: {e}")
         return False
 
-# ================== ИСПРАВЛЕННЫЙ ПРОМТ ДЛЯ УЛУЧШЕНИЯ ==================
+# ================== ПРАВИЛЬНЫЕ ПРОМТ-ШАБЛОНЫ ==================
 def get_enhancement_prompt(user_input: str) -> str:
-    """Промт, который заставляет модель генерировать детали, а не историю"""
+    """Создаем ПРАВИЛЬНЫЙ промт для улучшения, а не для рассказа"""
     return f"""<|system|>
-Ты помощник для создания промтов к AI-генератору изображений.
-Задача: взять короткое описание и добавить только визуальные детали.
+Ты помощник для создания промтов к AI-генератору изображений Dreamlike Photoreal 2.0.
+ТВОЯ ЗАДАЧА: Взять короткое описание и добавить КОНКРЕТНЫЕ детали для фотореалистичной генерации.
 
-Примеры как делать:
-"девушка" → "Молодая девушка с рыжими волосами и зелеными глазами, улыбка, мягкие черты лица"
-"лес" → "Магический лес с высокими деревьями, густая листва, солнечные лучи пробиваются сквозь ветви"
-"ночь" → "Темное ночное небо, яркие звезды, лунный свет, тихая атмосфера"
+ДОБАВЬ:
+1. Описание внешности (если есть персонаж)
+2. Детали окружения
+3. Освещение и время суток
+4. Атмосферу и настроение
+5. Стиль изображения
 
-НЕ ПИШИ: что происходит, что делает персонаж, сюжет, историю, действия
-ПИШИ ТОЛЬКО: как выглядит, какие детали, какое освещение, какая атмосфера
+НЕ пиши рассказ, НЕ добавляй сюжет. Только описание для генератора изображений.
 
-Формат ответа: только описание через запятую, заканчивая словами "фотореалистично, высокое качество, детализированно"</s>
+Примеры правильных улучшений:
+• "девушка в лесу" → "Молодая девушка с рыжими волосами в белом платье стоит в магическом лесу на закате, золотые лучи солнца пробиваются сквозь листву, атмосферное освещение, фотореалистично, детализированно"
+• "город ночью" → "Футуристический город ночью с неоновыми вывесками, мокрые улицы отражают свет, киберпанк стиль, детализированная архитектура, фотореалистично"</s>
 <|user|>
-Улучши это описание: "{user_input}"</s>
+Улучши это описание для генерации фотореалистичного изображения: "{user_input}"</s>
 <|assistant|>
-"""
+Улучшенное описание:"""
 
-# ================== ИСПРАВЛЕННЫЙ НЕГАТИВНЫЙ ПРОМТ ==================
-def generate_smart_negative_prompt(english_prompt: str, russian_prompt: str) -> str:
-    """Генерирует УМНЫЙ негативный промт на основе контекста"""
-    
-    # Базовые всегда исключаем
-    base_negative = [
-        "blurry", "ugly", "deformed", "bad anatomy",
-        "bad proportions", "extra limbs", "mutated hands",
-        "poorly drawn hands", "text", "watermark", "signature",
-        "logo", "username", "worst quality", "low quality",
-        "jpeg artifacts", "out of focus", "grainy", "noisy"
-    ]
-    
-    # Анализируем промты
-    combined_text = (english_prompt + " " + russian_prompt).lower()
-    
-    # Определяем что ИСКЛЮЧАТЬ на основе того что ХОТИМ
-    if any(word in combined_text for word in ["реалистич", "реалистично", "фото", "photoreal", "realistic"]):
-        # Если хотим реалистичное - исключаем нереалистичное
-        base_negative.extend(["cartoon", "anime", "painting", "drawing", "illustration", 
-                            "3d render", "digital art", "vector art", "sketch"])
-    
-    if any(word in combined_text for word in ["ярк", "свет", "солн", "bright", "light", "sun"]):
-        # Если хотим яркое - исключаем темное
-        base_negative.extend(["dark", "dull", "underexposed", "gloomy", "shadowy", "poor lighting"])
-    
-    if any(word in combined_text for word in ["детал", "подроб", "detailed", "intricate"]):
-        # Если хотим детализированное - исключаем простое
-        base_negative.extend(["simple", "plain", "boring", "empty", "basic"])
-    
-    if any(word in combined_text for word in ["красив", "прекрасн", "beautiful", "gorgeous"]):
-        # Если хотим красивое - исключаем некрасивое
-        base_negative.extend(["unattractive", "messy", "dirty", "disgusting"])
-    
-    if any(word in combined_text for word in ["портрет", "лицо", "человек", "portrait", "face", "person"]):
-        # Если есть лица - исключаем дефекты лиц
-        base_negative.extend(["asymmetric eyes", "cloned face", "malformed face", "bad eyes"])
-    
-    if any(word in combined_text for word in ["природ", "лес", "дерев", "nature", "forest", "tree"]):
-        # Если природа - исключаем искусственное
-        base_negative.extend(["artificial", "plastic", "synthetic", "man-made"])
-    
-    if any(word in combined_text for word in ["город", "здан", "архитект", "city", "building", "architecture"]):
-        # Если архитектура - исключаем разрушенное
-        base_negative.extend(["ruined", "destroyed", "broken", "crumbling"])
-    
-    # Убираем дубликаты и ограничиваем длину
-    unique_negative = []
-    for term in base_negative:
-        if term not in unique_negative:
-            unique_negative.append(term)
-    
-    # Берем первые 25 терминов (чтоб не слишком длинно)
-    return ", ".join(unique_negative[:25])
-
-def enhance_prompt_correct(user_input: str) -> str:
-    """Правильное улучшение промта"""
-    if _model is None:
-        return user_input + ", фотореалистично, высокое качество, детализированно"
+# ================== БЫСТРАЯ ГЕНЕРАЦИЯ ==================
+def enhance_prompt_fast(user_input: str) -> str:
+    """Быстро улучшаем промт с TinyLlama"""
+    if _model is None or _tokenizer is None:
+        return user_input + " [система не загружена]"
     
     try:
+        # Создаем правильный промт
         system_prompt = get_enhancement_prompt(user_input)
         
+        # Токенизация (ограничиваем длину)
         inputs = _tokenizer(
             system_prompt,
             return_tensors="pt",
-            max_length=300,
-            truncation=True
+            max_length=256,
+            truncation=True,
+            padding=True
         )
         
+        # Быстрая генерация с оптимизированными параметрами
         with torch.no_grad():
             outputs = _model.generate(
                 **inputs,
-                max_new_tokens=100,
-                temperature=0.7,
+                max_new_tokens=120,  # Не больше 120 новых токенов
+                temperature=0.8,
                 do_sample=True,
                 top_p=0.9,
+                repetition_penalty=1.1,
                 pad_token_id=_tokenizer.eos_token_id,
-                eos_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)
         
-        # Извлекаем только ответ ассистента
-        if "<|assistant|>" in response:
-            parts = response.split("<|assistant|>")
-            if len(parts) > 1:
-                enhanced = parts[1].strip()
-            else:
-                enhanced = response.replace(system_prompt, "").strip()
+        # Извлекаем только улучшенную часть
+        if "Улучшенное описание:" in response:
+            enhanced = response.split("Улучшенное описание:")[-1].strip()
         else:
+            # Fallback: удаляем системный промт
             enhanced = response.replace(system_prompt, "").strip()
         
-        # Очищаем
-        for marker in ["<|endoftext|>", "</s>", "<|user|>", "<|system|>"]:
+        # Очищаем от метаданных
+        for marker in ["<|endoftext|>", "</s>", "<|assistant|>", "<|user|>", "<|system|>"]:
             enhanced = enhanced.split(marker)[0].strip()
         
-        # Если модель написала историю - исправляем
-        if any(word in enhanced.lower() for word in ["который", "которая", "когда", "потом", "затем", "начала", "стал"]):
-            # Модель написала историю - используем fallback
-            enhanced = user_input + ", фотореалистично, высокое качество, детализированно, красивое освещение, детализированный фон"
+        # Если результат слишком короткий, добавляем базовые улучшения
+        if len(enhanced.split()) < 5:
+            enhanced = f"{user_input}, фотореалистично, высокое качество, детализированно, красивое освещение"
         
-        # Добавляем технические термины если их нет
-        if "фотореалистично" not in enhanced.lower():
-            enhanced += ", фотореалистично"
-        if "высокое качество" not in enhanced.lower():
-            enhanced += ", высокое качество"
-        if "детализированно" not in enhanced.lower():
-            enhanced += ", детализированно"
-        
-        return enhanced[:350]
+        return enhanced[:400]  # Ограничиваем длину
         
     except Exception as e:
         print(f"Ошибка генерации: {e}")
-        return user_input + ", фотореалистично, высокое качество, детализированно"
+        return f"{user_input}, фотореалистично, высокое качество, детализированно"
 
-def process_pipeline_full(user_input: str):
-    """Полный пайплайн с УМНЫМ негативным промтом"""
-    if not user_input.strip():
-        return "Введите промт", "", "", "⏳ Введите промт"
+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"
+    ]
+    
+    prompt_lower = english_prompt.lower()
+    
+    # Добавляем контекстные исключения
+    if any(word in prompt_lower for word in ["realistic", "photoreal", "photo"]):
+        base.extend(["cartoon", "anime", "painting", "drawing", "illustration", "3d render"])
+    
+    if any(word in prompt_lower for word in ["bright", "light", "sun"]):
+        base.extend(["dark", "dull", "underexposed", "gloomy"])
     
-    print(f"🧠 Обработка: {user_input}")
+    if any(word in prompt_lower for word in ["portrait", "face", "person"]):
+        base.extend(["asymmetric eyes", "cloned face", "malformed limbs"])
     
+    # Убираем дубликаты
+    unique = []
+    for term in base:
+        if term not in unique:
+            unique.append(term)
+    
+    return ", ".join(unique)
+
+def process_pipeline(user_input: str):
+    """Полный пайплайн обработки"""
     start_time = time.time()
     
-    # 1. Улучшаем русский
+    if not user_input.strip():
+        return "Введите промт", "", "", "⏳ Введите промт"
+    
+    print(f"⚡ Обработка: {user_input}")
+    
+    # 1. Улучшаем промт (максимум 5 секунд)
     gen_start = time.time()
-    enhanced_ru = enhance_prompt_correct(user_input)
+    enhanced_ru = enhance_prompt_fast(user_input)
     gen_time = time.time() - gen_start
     
     # 2. Переводим
     trans_start = time.time()
     try:
         if _translator:
-            translated = _translator(enhanced_ru[:350])[0]['translation_text']
+            translated = _translator(enhanced_ru[:300])[0]['translation_text']
         else:
             translated = enhanced_ru
     except:
         translated = enhanced_ru
     trans_time = time.time() - trans_start
     
-    # 3. Улучшаем английский
-    dreamlike_keywords = ", photorealistic, high quality, detailed, 8K, masterpiece"
-    keywords_to_add = []
-    
+    # 3. Улучшаем английский для Dreamlike
+    dreamlike_keywords = ", photorealistic, high quality, detailed, 8K, professional photography"
     if "photorealistic" not in translated.lower():
-        keywords_to_add.append("photorealistic")
-    if "high quality" not in translated.lower():
-        keywords_to_add.append("high quality")
-    if "detailed" not in translated.lower():
-        keywords_to_add.append("detailed")
-    if "8k" not in translated.lower():
-        keywords_to_add.append("8K")
-    
-    if keywords_to_add:
-        translated += ", " + ", ".join(keywords_to_add)
+        translated += dreamlike_keywords
     
-    # 4. Генерируем УМНЫЙ негативный промт
-    neg_start = time.time()
-    negative_en = generate_smart_negative_prompt(translated, enhanced_ru)
-    neg_time = time.time() - neg_start
+    # 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"""
-    ⚡ СКОРОСТЬ ОБРАБОТКИ:
-    • Улучшение: {gen_time:.1f} сек
-    • Перевод: {trans_time:.1f} сек  
-    • Негативный промт: {neg_time:.1f} сек
+    ⚡ БЫСТРАЯ ОБРАБОТКА:
+    • Генерация: {gen_time:.1f} сек
+    • Перевод: {trans_time:.1f} сек
+    • Негативный промт: {negative_time:.1f} сек
     • Всего: {total_time:.1f} сек
     
-    🎨 КАЧЕСТВО:
-    • Улучшенный промт: {len(enhanced_ru.split())} слов
-    • Английский промт: {len(translated.split())} слов
-    • Негативных терминов: {len(negative_en.split(','))}
-    • Статус: ✅ Готов для Dreamlike Photoreal 2.0
+    📊 СТАТИСТИКА:
+    • Модель: TinyLlama-1.1B (оптимизирована для CPU)
+    • Длина промта: {len(translated.split())} слов
+    • Статус: ✅ Готово к использованию в Dreamlike
     """
     
     return enhanced_ru, translated, negative_en, info.strip()
 
-# ================== ИНТЕРФЕЙС ==================
-with gr.Blocks(title="Умный улучшитель промтов", theme=gr.themes.Soft()) as demo:
+# ================== ПРОСТОЙ ИНТЕРФЕЙС ==================
+with gr.Blocks(title="AI Улучшитель промтов для Dreamlike", theme=gr.themes.Soft()) as demo:
     
+    # Заголовок
     gr.Markdown("""
-    # 🧠 УМНЫЙ УЛУЧШИТЕЛЬ ПРОМТОВ
-    ### AI улучшает • Умный негативный промт • Для Dreamlike Photoreal
+    # 🚀 УЛЬТРА-БЫСТРЫЙ AI УЛУЧШИТЕЛЬ ПРОМТОВ
+    ### Оптимизирован для CPU • TinyLlama • Мгновенная обработка
     """)
     
-    # Статус
+    # Статус загрузки
     status = gr.Textbox(
-        label="🔄 Статус",
+        label="🔄 Статус системы",
         value="Загружаю систему...",
         interactive=False
     )
@@ -267,70 +235,90 @@ with gr.Blocks(title="Умный улучшитель промтов", theme=gr.
         # Ввод
         input_prompt = gr.Textbox(
             label="🎯 Введите промт на русском",
-            placeholder="девушка в магическом лесу на закате",
+            placeholder="Пример: девушка в магическом лесу на закате",
             lines=2
         )
         
         # Кнопка
-        btn = gr.Button("🧠 УЛУЧШИТЬ С AI", variant="primary", size="lg")
+        process_btn = gr.Button("⚡ УЛУЧШИТЬ ПРОМТ", variant="primary", size="lg")
         
         # Результаты
-        gr.Markdown("### 📝 УЛУЧШЕННЫЙ (AI генерация)")
+        gr.Markdown("### 📝 УЛУЧШЕННЫЙ ПРОМТ (русский)")
         output_ru = gr.Textbox(
+            label="",
             lines=3,
-            show_label=False
+            interactive=True
         )
         
-        gr.Markdown("### 🌐 ДЛЯ DREAMLIKE")
+        gr.Markdown("### 🌐 ДЛЯ DREAMLIKE (английский)")
         with gr.Row():
             output_en = gr.Textbox(
                 label="🎨 Основной промт",
                 lines=3,
                 interactive=True
             )
+            
             output_neg = gr.Textbox(
-                label="🚫 УМНЫЙ Negative Prompt",
+                label="🚫 Negative Prompt",
                 lines=3,
                 interactive=True
             )
         
-        # Инфо
+        # Информация
         output_info = gr.Textbox(
-            label="📊 СТАТИСТИКА И ИНФО",
-            lines=4
+            label="📊 СТАТИСТИКА",
+            lines=4,
+            interactive=False
         )
         
         # Примеры
-        gr.Markdown("### 🧪 ТЕСТ-КЕЙСЫ")
+        gr.Markdown("### 🧪 БЫСТРЫЕ ПРИМЕРЫ")
         examples = gr.Examples(
             examples=[
-                ["девушка в лесу"],
-                ["космонавт на луне"],
-                ["киберпанк город ночью"],
-                ["дракон на горе"],
-                ["старый волшебник"]
+                ["девушка в магическом лесу на закате"],
+                ["киберпанк улица ночью с неоновыми вывесками"],
+                ["космонавт играет на гитаре на луне"],
+                ["кошка в шляпе читает газету в кафе"],
+                ["дракон на вершине горы"]
             ],
             inputs=[input_prompt],
             outputs=[output_ru, output_en, output_neg, output_info],
-            fn=process_pipeline_full,
+            fn=process_pipeline,
             label="Кликните для теста:"
         )
         
         # Обработчики
-        def process(prompt):
-            return process_pipeline_full(prompt)
+        def process_and_update(prompt):
+            return process_pipeline(prompt)
+        
+        process_btn.click(
+            fn=process_and_update,
+            inputs=[input_prompt],
+            outputs=[output_ru, output_en, output_neg, output_info]
+        )
         
-        btn.click(fn=process, inputs=[input_prompt], outputs=[output_ru, output_en, output_neg, output_info])
-        input_prompt.submit(fn=process, 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]
+        )
     
-    # Инициализация
-    def init():
+    # Функция инициализации
+    def initialize_system():
         if load_system():
-            return "✅ СИСТЕМА ГОТОВА! Введите промт.", gr.update(visible=True)
-        return "❌ Ошибка загрузки", gr.update(visible=False)
+            return "✅ СИСТЕМА ГОТОВА! TinyLlama загружена и оптимизирована.\n⚡ Введите промт и нажмите кнопку.", gr.update(visible=True)
+        else:
+            return "❌ Ошибка загрузки. Попробуйте перезапустить Space.", gr.update(visible=False)
     
-    demo.load(fn=init, outputs=[status, main_ui])
+    # Загружаем при старте
+    demo.load(
+        fn=initialize_system,
+        outputs=[status, main_ui]
+    )
 
-# Запуск
+# ================== ЗАПУСК ==================
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860
+    )