app.py

import streamlit as st
import torch
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
from midiutil import MIDIFile
from PIL import Image, ImageFilter, ImageEnhance
import io
import hashlib
import time
import random
import numpy as np

# Настройки страницы
st.set_page_config(
    page_title="🎨 Pro AI Генератор",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# CSS стили
st.markdown("""
<style>
    .main-header {
        text-align: center;
        font-size: 3rem;
        background: linear-gradient(45deg, #FF6B6B, #4ECDC4, #45B7D1);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        margin-bottom: 1rem;
    }
    .feature-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin: 10px 0;
    }
    .stButton > button {
        background: linear-gradient(45deg, #FF6B6B, #FF8E53);
        color: white;
        font-weight: bold;
        border-radius: 10px;
        border: none;
        padding: 12px 24px;
        transition: all 0.3s;
    }
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 5px 15px rgba(255, 107, 107, 0.4);
    }
</style>
""", unsafe_allow_html=True)

# Заголовок
st.markdown('<h1 class="main-header">🚀 Pro AI Генератор</h1>', unsafe_allow_html=True)
st.markdown("### 🎨 Изображения • 🎵 Музыка • ⚡ Продвинутые функции")

# ====================
# ИНИЦИАЛИЗАЦИЯ
# ====================

# Инициализация состояния
if 'history' not in st.session_state:
    st.session_state.history = []
if 'model' not in st.session_state:
    st.session_state.model = None

# ====================
# ФУНКЦИИ ДЛЯ МУЗЫКИ
# ====================

def create_chord_progression(key=60, progression_type="basic"):
    """Создание аккордовых прогрессий"""
    progressions = {
        "basic": [[0, 4, 7], [5, 9, 0], [7, 11, 2], [2, 5, 9]],  # I-V-vi-IV
        "jazz": [[0, 4, 7, 10], [2, 5, 9, 0], [7, 11, 2, 5], [9, 0, 4, 7]],  # ii-V-I-vi
        "blues": [[0, 3, 7], [0, 3, 7], [4, 7, 10], [5, 8, 0]],  # Блюзовая
        "epic": [[0, 4, 7], [4, 8, 11], [5, 9, 0], [2, 5, 9]],  # Эпичная
    }
    return [[key + note for note in chord] for chord in progressions.get(progression_type, progressions["basic"])]

def safe_encode(text):
    """Безопасное кодирование текста"""
    try:
        return text.encode('utf-8')
    except:
        return str(text).encode('utf-8', errors='ignore')

def generate_advanced_music(prompt, duration=15, style="Веселая", 
                           tempo=120, complexity="medium", use_chords=True):
    """Продвинутая генерация музыки"""
    try:
        midi = MIDIFile(2)  # 2 трека: мелодия и аккомпанемент
        time_pos = 0
        
        # Настройки темпа
        midi.addTrackName(0, time_pos, "Melody")
        midi.addTrackName(1, time_pos, "Accompaniment")
        midi.addTempo(0, time_pos, tempo)
        midi.addTempo(1, time_pos, tempo)
        
        # Определяем гамму и тональность по стилю
        prompt_lower = str(prompt).lower()
        
        if "sad" in prompt_lower or "груст" in prompt_lower:
            key = 60  # C минор
            scale = [0, 2, 3, 5, 7, 8, 10]  # Минорная гамма
            progression = "jazz"
            tempo = min(tempo, 100)  # Медленнее для грустной
        elif "epic" in prompt_lower or "эпич" in prompt_lower:
            key = 60
            scale = [0, 2, 4, 5, 7, 9, 11]  # Мажорная гамма
            progression = "epic"
            tempo = max(tempo, 140)  # Быстрее для эпичной
        elif "mysterious" in prompt_lower or "таинств" in prompt_lower:
            key = 60
            scale = [0, 1, 4, 5, 7, 8, 11]  # Мистическая гамма
            progression = "jazz"
        elif "chill" in prompt_lower or "расслаб" in prompt_lower:
            key = 60
            scale = [0, 2, 3, 5, 7, 9, 10]
            progression = "basic"
            tempo = 100
        else:  # Веселая по умолчанию
            key = 60  # C мажор
            scale = [0, 2, 4, 5, 7, 9, 11]  # Мажорная гамма
            progression = "basic"
        
        # Безопасное создание хэша
        try:
            prompt_bytes = safe_encode(prompt)
            hash_bytes = hashlib.md5(prompt_bytes).digest()
        except:
            hash_bytes = hashlib.md5(b"default").digest()
        
        # Создаем аккордовую прогрессию
        chords = create_chord_progression(key, progression)
        
        # Генерация аккомпанемента (аккорды)
        if use_chords:
            chord_duration = 2.0 if complexity == "simple" else 1.0
            max_chords = min(len(chords), int(duration / chord_duration))
            
            for i in range(max_chords):
                chord_idx = i % len(chords)
                for note_offset in chords[chord_idx]:
                    midi.addNote(1, 0, note_offset, 
                               time_pos + i * chord_duration, 
                               chord_duration, 60)
        
        # Генерация мелодии
        notes_per_beat = 2 if complexity == "complex" else 1
        total_beats = duration * (tempo / 60)  # Количество долей
        total_notes = min(int(total_beats * notes_per_beat), 500)  # Ограничиваем
        
        for beat in range(total_notes):
            # Используем хэш для детерминированной генерации
            if len(hash_bytes) > 0:
                hash_val = hash_bytes[beat % len(hash_bytes)]
            else:
                hash_val = beat % 256
            
            # Безопасный доступ к scale
            if scale and len(scale) > 0:
                note_idx = hash_val % len(scale)
                note = key + scale[note_idx]
            else:
                note = 60 + (hash_val % 12)  # Fallback: хроматическая гамма
            
            # Добавляем октавные вариации
            octave = (hash_val // 7) % 3 - 1  # 7 нот в гамме
            note += octave * 12
            
            # Ограничиваем диапазон MIDI (21-108)
            note = max(21, min(108, note))
            
            # Динамика (громкость)
            velocity = 70 + (hash_val % 30)
            
            # Длительность ноты
            note_duration = 0.5 / notes_per_beat
            
            # Ритмический паттерн
            if hash_val % 8 == 0 and complexity != "simple":
                note_duration *= 2  # Более длинные ноты иногда
            
            # Добавляем ноту
            midi.addNote(0, 0, int(note), 
                        time_pos + beat * (1.0 / notes_per_beat), 
                        note_duration, velocity)
        
        # Добавляем басовую линию
        if complexity == "complex" and chords and len(chords) > 0:
            max_bass_notes = min(int(duration), 100)
            for i in range(max_bass_notes):
                chord_idx = i % len(chords)
                if chords[chord_idx]:
                    bass_note = chords[chord_idx][0] - 12  # На октаву ниже
                    bass_note = max(21, min(108, bass_note))
                    midi.addNote(1, 0, int(bass_note), 
                               time_pos + i * 1.0, 1.0, 50)
        
        # Сохраняем
        with io.BytesIO() as output:
            midi.writeFile(output)
            return output.getvalue()
            
    except Exception as e:
        st.error(f"Ошибка генерации музыки: {str(e)}")
        return None

# ====================
# ФУНКЦИИ ДЛЯ ИЗОБРАЖЕНИЙ
# ====================

@st.cache_resource
def load_advanced_model(model_choice="small"):
    """Загрузка разных моделей Stable Diffusion"""
    model_map = {
        "small": "OFA-Sys/small-stable-diffusion-v0",
        "standard": "runwayml/stable-diffusion-v1-5",
    }
    
    try:
        model_id = model_map.get(model_choice, model_map["small"])
        
        # Загружаем с разными планировщиками для качества
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float32,
            safety_checker=None,
            requires_safety_checker=False
        )
        
        # Используем Euler для лучшего качества
        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
        
        return pipe
    except Exception as e:
        st.error(f"Ошибка загрузки модели: {str(e)}")
        return None

def generate_advanced_image(prompt, style="Реализм", 
                           size=(512, 512), steps=30, guidance=7.5,
                           negative_prompt="", model_choice="small"):
    """Продвинутая генерация изображений"""
    try:
        # Добавляем стилевые ключевые слова
        style_prompts = {
            "Реализм": "photorealistic, detailed, 8k",
            "Аниме": "anime style, vibrant colors, detailed",
            "Цифровое искусство": "digital art, concept art, trending on artstation",
            "Масляная живопись": "oil painting, brush strokes, masterpiece",
            "Пиксель-арт": "pixel art, 8-bit, retro game style",
        }
        
        enhanced_prompt = f"{prompt}, {style_prompts.get(style, '')}"
        
        # Загружаем выбранную модель
        model = load_advanced_model(model_choice)
        if model is None:
            return None, "Не удалось загрузить модель"
        
        # Генерация
        with torch.no_grad():
            result = model(
                prompt=enhanced_prompt,
                negative_prompt=negative_prompt,
                width=size[0],
                height=size[1],
                num_inference_steps=steps,
                guidance_scale=guidance,
                num_images_per_prompt=1
            )
        
        image = result.images[0]
        
        return image, "Успешно"
        
    except Exception as e:
        return None, f"Ошибка: {str(e)}"

# ====================
# ИНТЕРФЕЙС
# ====================

# Боковая панель
with st.sidebar:
    st.markdown('<div class="feature-card">', unsafe_allow_html=True)
    st.markdown("### ⚡ Продвинутые функции")
    st.markdown("""
    - 🎵 Умная генерация музыки
    - 🎨 Стилизованные изображения
    - 🎭 Аккордовые прогрессии
    - 📊 История генераций
    """)
    st.markdown('</div>', unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Настройки
    mode = st.radio(
        "🎯 Режим работы:",
        ["Музыка", "Изображения"],
        index=0
    )

# Основной контент
if mode == "Музыка":
    st.header("🎵 Продвинутый генератор музыки")
    
    # Колонки для настроек
    col1, col2 = st.columns(2)
    
    with col1:
        music_prompt = st.text_area(
            "📝 Описание музыки:",
            "Epic adventure theme",
            height=100
        )
        
        # Выбор стиля
        music_style = st.selectbox(
            "🎭 Музыкальный стиль:",
            ["Веселая", "Грустная", "Эпичная", "Таинственная", "Расслабляющая"],
            index=0
        )
        
        # Сложность
        complexity = st.select_slider(
            "📊 Сложность композиции:",
            options=["simple", "medium", "complex"],
            value="medium"
        )
    
    with col2:
        duration = st.slider("⏱️ Длительность (секунд):", 10, 60, 30)
        tempo = st.slider("🎼 Темп (BPM):", 60, 200, 120)
        
        # Дополнительные опции
        use_chords = st.checkbox("🎹 Добавить аккорды", value=True)
    
    # Кнопка генерации
    if st.button("🚀 Создать продвинутую музыку", type="primary", use_container_width=True):
        if not music_prompt:
            st.warning("Введите описание музыки!")
        else:
            with st.spinner(f"🎶 Создание {music_style} музыки..."):
                start_time = time.time()
                
                data = generate_advanced_music(
                    prompt=music_prompt,
                    duration=duration,
                    style=music_style,
                    tempo=tempo,
                    complexity=complexity,
                    use_chords=use_chords
                )
                
                if data is not None:
                    elapsed = time.time() - start_time
                    
                    st.session_state.history.append({
                        "type": "music",
                        "content": data,
                        "prompt": music_prompt,
                        "style": music_style,
                        "time": time.strftime("%H:%M:%S"),
                        "duration": f"{duration} сек",
                        "tempo": f"{tempo} BPM"
                    })
                    
                    st.success(f"✅ Музыка создана за {elapsed:.1f} сек!")
                    
                    # Показываем информацию о треке
                    col_info1, col_info2 = st.columns(2)
                    with col_info1:
                        st.info(f"**Стиль:** {music_style}")
                        st.info(f"**Темп:** {tempo} BPM")
                    with col_info2:
                        st.info(f"**Длительность:** {duration} сек")
                        st.info(f"**Сложность:** {complexity}")
                    
                    # Скачивание
                    filename = f"advanced_music_{int(time.time())}.mid"
                    st.download_button(
                        "📥 Скачать MIDI файл",
                        data=data,
                        file_name=filename,
                        mime="audio/midi",
                        use_container_width=True
                    )
                else:
                    st.error("Не удалось создать музыку")

else:  # Изображения
    st.header("🎨 Продвинутый генератор изображений")
    
    # Три колонки для настроек
    col1, col2, col3 = st.columns(3)
    
    with col1:
        image_prompt = st.text_area(
            "📝 Описание изображения:",
            "A majestic dragon flying over ancient mountains",
            height=120
        )
        
        # Выбор стиля
        image_style = st.selectbox(
            "🎨 Стиль изображения:",
            ["Реализм", "Аниме", "Цифровое искусство", "Масляная живопись", "Пиксель-арт"],
            index=0
        )
        
        # Выбор модели
        model_choice = st.selectbox(
            "🤖 Модель генерации:",
            ["small", "standard"],
            index=0,
            help="small - быстрее, standard - качественнее"
        )
    
    with col2:
        # Размеры
        size_option = st.selectbox(
            "📏 Размер изображения:",
            ["512x512", "768x512", "512x768", "768x768"],
            index=0
        )
        width, height = map(int, size_option.split('x'))
        
        # Качество
        steps = st.slider("🎯 Качество (шагов):", 20, 50, 30)
        guidance = st.slider("📐 Строгость промпта:", 5.0, 15.0, 7.5)
    
    with col3:
        # Негативный промпт
        negative_prompt = st.text_area(
            "🚫 Исключить из изображения:",
            "blurry, ugly, deformed, poorly drawn",
            height=100
        )
    
    # Кнопка генерации
    if st.button("🚀 Сгенерировать продвинутое изображение", type="primary", use_container_width=True):
        if not image_prompt:
            st.warning("Введите описание изображения!")
        else:
            with st.spinner(f"🎨 Генерация {image_style} изображения..."):
                start_time = time.time()
                
                image, status = generate_advanced_image(
                    prompt=image_prompt,
                    style=image_style,
                    size=(width, height),
                    steps=steps,
                    guidance=guidance,
                    negative_prompt=negative_prompt,
                    model_choice=model_choice
                )
                
                if image is not None:
                    elapsed = time.time() - start_time
                    
                    st.session_state.history.append({
                        "type": "image",
                        "content": image,
                        "prompt": image_prompt,
                        "style": image_style,
                        "time": time.strftime("%H:%M:%S"),
                        "size": size_option,
                        "quality": f"{steps} шагов"
                    })
                    
                    st.success(f"✅ Изображение создано за {elapsed:.1f} сек!")
                    
                    # Показываем изображение
                    st.image(image, use_column_width=True, 
                            caption=f"Стиль: {image_style} | Размер: {size_option}")
                    
                    # Информация
                    col_info1, col_info2 = st.columns(2)
                    with col_info1:
                        st.info(f"**Стиль:** {image_style}")
                        st.info(f"**Модель:** {model_choice}")
                    with col_info2:
                        st.info(f"**Размер:** {size_option}")
                        st.info(f"**Качество:** {steps} шагов")
                    
                    # Скачивание
                    buf = io.BytesIO()
                    image.save(buf, format="PNG", optimize=True)
                    filename = f"advanced_image_{int(time.time())}.png"
                    
                    st.download_button(
                        "📥 Скачать PNG",
                        data=buf.getvalue(),
                        file_name=filename,
                        mime="image/png",
                        use_container_width=True
                    )
                else:
                    st.error(f"❌ {status}")

# ====================
# ИСТОРИЯ ГЕНЕРАЦИЙ
# ====================

if st.session_state.history:
    st.markdown("---")
    st.header("📊 История генераций")
    
    # Показываем историю (последние 5)
    for i, item in enumerate(reversed(st.session_state.history[-5:])):
        with st.expander(f"{'🎵' if item['type'] == 'music' else '🎨'} "
                        f"Генерация {len(st.session_state.history)-i} - {item['time']}", 
                        expanded=i==0):
            
            if item["type"] == "music":
                col1, col2, col3 = st.columns([3, 1, 1])
                with col1:
                    st.write(f"**{item['prompt'][:50]}...**")
                    st.caption(f"Стиль: {item.get('style', 'Не указан')}")
                    st.caption(f"Длительность: {item.get('duration', 'Не указано')}")
                with col2:
                    st.download_button(
                        "🎵 Скачать",
                        data=item["content"],
                        file_name=f"history_music_{i}.mid",
                        mime="audio/midi",
                        key=f"hist_mus_{i}"
                    )
                with col3:
                    if st.button("🗑️", key=f"del_mus_{i}"):
                        st.session_state.history.remove(item)
                        st.rerun()
            
            else:  # image
                col1, col2, col3 = st.columns([3, 1, 1])
                with col1:
                    st.image(item["content"], width=200)
                    st.caption(f"**{item['prompt'][:30]}...**")
                    st.caption(f"Стиль: {item.get('style', 'Не указан')}")
                with col2:
                    buf = io.BytesIO()
                    item["content"].save(buf, format="PNG")
                    st.download_button(
                        "💾 Скачать",
                        data=buf.getvalue(),
                        file_name=f"history_img_{i}.png",
                        mime="image/png",
                        key=f"hist_img_{i}"
                    )
                with col3:
                    if st.button("🗑️", key=f"del_img_{i}"):
                        st.session_state.history.remove(item)
                        st.rerun()

# ====================
# ФУТЕР
# ====================

st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666; padding: 20px;'>
<h3>🚀 Pro AI Генератор</h3>
<p>🎵 Умная генерация музыки | 🎨 Стилизованные изображения | ⚡ Продвинутые функции</p>
<p>🤗 Работает на Hugging Face Spaces | 🐍 Python 3.10</p>
</div>
""", unsafe_allow_html=True)

# Статистика
with st.expander("📈 Статистика", expanded=False):
    col_stat1, col_stat2 = st.columns(2)
    with col_stat1:
        st.metric("Всего генераций", len(st.session_state.history))
    with col_stat2:
        music_count = len([h for h in st.session_state.history if h["type"] == "music"])
        st.metric("Музыкальных треков", music_count)