Update app.py

app.py

@@ -7,37 +7,32 @@ from PIL import Image
 import imageio
 import os
 import tempfile
-from datetime import datetime
 
-# ============ ПОЛНОСТЬЮ НЕЙРОСЕТЕВАЯ АРХИТЕКТУРА ============
+# ============ ИСПРАВЛЕННАЯ ПОЛНОСТЬЮ НЕЙРОСЕТЕВАЯ АРХИТЕКТУРА ============
 class FullNeuralAnimator(nn.Module):
     """
-    Одна нейросеть делает ВСЁ:
-    1. Анализирует изображение
-    2. Предсказывает последовательность кадров
-    3. Генерирует анимацию
+    Одна нейросеть делает ВСЁ
     """
     def __init__(self):
         super().__init__()
         
-        # === Encoder (понимает структуру) ===
+        # === Encoder ===
         self.enc1 = self._block(3, 32)
         self.enc2 = self._block(32, 64)
         self.enc3 = self._block(64, 128)
         self.enc4 = self._block(128, 256)
         self.pool = nn.MaxPool2d(2)
         
-        # === LSTM для временной последовательности ===
-        # Запоминает как меняется анимация во времени
+        # === LSTM ===
         self.lstm = nn.LSTM(
-            input_size=256 * 16 * 16,  # 256 каналов * 16x16
+            input_size=256 * 8 * 8,  # 256 каналов * 8x8 (после 3х пулингов)
             hidden_size=512,
             num_layers=2,
             batch_first=True,
             dropout=0.2
         )
         
-        # === Декодер (создаёт кадры) ===
+        # === Декодер ===
         self.dec4 = self._block(512, 256)
         self.dec3 = self._block(256, 128)
         self.dec2 = self._block(128, 64)
@@ -56,9 +51,6 @@ class FullNeuralAnimator(nn.Module):
             nn.Tanh()
         )
         
-        # === Контроль времени ===
-        self.time_encoder = nn.Linear(1, 128)  # Кодируем время
-        
     def _block(self, in_ch, out_ch):
         return nn.Sequential(
             nn.Conv2d(in_ch, out_ch, 3, padding=1),
@@ -70,10 +62,6 @@ class FullNeuralAnimator(nn.Module):
         )
     
     def forward(self, x, num_frames=20):
-        """
-        x: входное изображение [B, 3, H, W]
-        num_frames: сколько кадров сгенерировать
-        """
         batch_size = x.size(0)
         
         # === 1. Кодируем изображение ===
@@ -86,45 +74,45 @@ class FullNeuralAnimator(nn.Module):
         skips = [e1, e2, e3, e4]
         
         # === 2. Подготовка для LSTM ===
-        # [B, 256, 16, 16] -> [B, 256*16*16]
-        bottleneck = e4.view(batch_size, -1)
+        # После 3х пулингов: 256 -> 8x8
+        bottleneck = e4.view(batch_size, -1)  # [B, 256*8*8]
         
-        # === 3. Генерируем последовательность во времени ===
+        # === 3. Генерируем последовательность ===
         frames = []
         hidden = None
         
-        # Начальное состояние
         lstm_input = bottleneck.unsqueeze(1)  # [B, 1, features]
         
         for t in range(num_frames):
-            # Кодируем время
-            time_tensor = torch.tensor([t / num_frames], device=x.device)
-            time_embed = self.time_encoder(time_tensor).unsqueeze(0).unsqueeze(1)  # [1, 1, 128]
-            
-            # Добавляем информацию о времени
-            lstm_input_with_time = torch.cat([lstm_input, time_embed.repeat(batch_size, 1, 1)], dim=-1)
-            
             # LSTM предсказывает следующее состояние
-            lstm_out, hidden = self.lstm(lstm_input_with_time, hidden)
+            lstm_out, hidden = self.lstm(lstm_input, hidden)
             
             # === 4. Декодируем в кадр ===
-            # [B, 512] -> [B, 256, 16, 16]
-            h = lstm_out.squeeze(1).view(batch_size, 256, 16, 16)
+            h = lstm_out.squeeze(1).view(batch_size, 256, 8, 8)
             
             # Декодер с skip connections
             d4 = self.up4(h)
+            # Resize skip connection если нужно
+            if d4.size(-1) != skips[3].size(-1):
+                skips[3] = F.interpolate(skips[3], size=d4.size(-2:), mode='bilinear')
             d4 = torch.cat([d4, skips[3]], dim=1)
             d4 = self.dec4(d4)
             
             d3 = self.up3(d4)
+            if d3.size(-1) != skips[2].size(-1):
+                skips[2] = F.interpolate(skips[2], size=d3.size(-2:), mode='bilinear')
             d3 = torch.cat([d3, skips[2]], dim=1)
             d3 = self.dec3(d3)
             
             d2 = self.up2(d3)
+            if d2.size(-1) != skips[1].size(-1):
+                skips[1] = F.interpolate(skips[1], size=d2.size(-2:), mode='bilinear')
             d2 = torch.cat([d2, skips[1]], dim=1)
             d2 = self.dec2(d2)
             
             d1 = self.up1(d2)
+            if d1.size(-1) != skips[0].size(-1):
+                skips[0] = F.interpolate(skips[0], size=d1.size(-2:), mode='bilinear')
             d1 = torch.cat([d1, skips[0]], dim=1)
             d1 = self.dec1(d1)
             
@@ -132,34 +120,28 @@ class FullNeuralAnimator(nn.Module):
             frame = self.frame_generator(d1)
             frames.append(frame)
             
-            # Обновляем вход для LSTM (авторегрессия)
-            # Берём bottleneck следующего кадра
-            next_bottleneck = self.enc4(self.pool(self.enc3(self.pool(self.enc2(self.pool(self.enc1(frame)))))))
-            next_bottleneck = next_bottleneck.view(batch_size, -1)
-            lstm_input = next_bottleneck.unsqueeze(1)
-        
-        # Собираем все кадры
-        return torch.stack(frames, dim=1)  # [B, T, 3, H, W]
+            # Обновляем вход для LSTM
+            next_features = self.enc4(self.pool(self.enc3(self.pool(self.enc2(self.pool(self.enc1(frame)))))))
+            next_features = next_features.view(batch_size, -1)
+            lstm_input = next_features.unsqueeze(1)
+        
+        return torch.stack(frames, dim=1)
 
-# ============ ЕЩЁ ОДНА НЕЙРОСЕТЬ ДЛЯ РАЗНООБРАЗИЯ ============
+# ============ СТИЛЕВАЯ НЕЙРОСЕТЬ ============
 class StyleTransferAnimator(nn.Module):
-    """
-    Генерирует разные стили анимации
-    """
     def __init__(self):
         super().__init__()
         
-        # Стили анимации (обучаемые векторы)
+        # Стили
         self.style_embeddings = nn.ParameterDict({
             'wave': nn.Parameter(torch.randn(64)),
             'pulse': nn.Parameter(torch.randn(64)),
             'glitch': nn.Parameter(torch.randn(64)),
             'melt': nn.Parameter(torch.randn(64)),
             'twist': nn.Parameter(torch.randn(64)),
-            'dream': nn.Parameter(torch.randn(64)),
         })
         
-        # Основная сеть
+        # Encoder
         self.encoder = nn.Sequential(
             nn.Conv2d(3, 32, 4, stride=2, padding=1),
             nn.ReLU(),
@@ -171,53 +153,58 @@ class StyleTransferAnimator(nn.Module):
             nn.ReLU(),
         )
         
-        # Генератор кадров с учётом стиля
+        # LSTM
+        self.lstm = nn.LSTM(
+            input_size=256 * 8 * 8,
+            hidden_size=512,
+            num_layers=2,
+            batch_first=True,
+            dropout=0.2
+        )
+        
+        # Decoder
         self.decoder = nn.Sequential(
-            nn.ConvTranspose2d(256 + 64, 128, 4, stride=2, padding=1),
+            nn.ConvTranspose2d(512 + 64, 256, 4, stride=2, padding=1),
             nn.ReLU(),
-            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
+            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),
             nn.ReLU(),
-            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
+            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
             nn.ReLU(),
-            nn.ConvTranspose2d(32, 3, 4, stride=2, padding=1),
+            nn.ConvTranspose2d(64, 3, 4, stride=2, padding=1),
             nn.Tanh()
         )
-        
-        # LSTM для времени
-        self.temporal_lstm = nn.LSTMCell(256, 512)
-        self.time_proj = nn.Linear(1, 128)
     
     def forward(self, x, style='wave', num_frames=20):
         batch_size = x.size(0)
         
-        # Кодируем изображение
-        features = self.encoder(x)  # [B, 256, 16, 16]
-        features_flat = features.view(batch_size, -1)  # [B, 256*16*16]
+        # Кодируем
+        features = self.encoder(x)  # [B, 256, 8, 8]
+        features_flat = features.view(batch_size, -1)
         
-        # Получаем стиль
-        style_vector = self.style_embeddings[style]  # [64]
-        style_vector = style_vector.unsqueeze(0).repeat(batch_size, 1)  # [B, 64]
+        # Стиль
+        style_vector = self.style_embeddings[style]
+        style_vector = style_vector.unsqueeze(0).repeat(batch_size, 1)
         
         frames = []
-        h = None
-        c = None
+        hidden = None
+        lstm_input = features_flat.unsqueeze(1)
         
         for t in range(num_frames):
-            # Время
-            t_norm = torch.tensor([t / num_frames], device=x.device)
-            t_embed = self.time_proj(t_norm).unsqueeze(0).repeat(batch_size, 1)
-            
             # LSTM
-            lstm_input = torch.cat([features_flat, t_embed, style_vector], dim=1)
-            h, c = self.temporal_lstm(lstm_input, (h, c))
+            lstm_out, hidden = self.lstm(lstm_input, hidden)
             
-            # Декодируем
-            h_reshaped = h.view(batch_size, 256, 16, 16)
-            style_reshaped = style_vector.view(batch_size, 64, 1, 1).repeat(1, 1, 16, 16)
-            decoder_input = torch.cat([h_reshaped, style_reshaped], dim=1)
+            # Декодируем со стилем
+            h = lstm_out.squeeze(1).view(batch_size, 256, 8, 8)
+            style_expanded = style_vector.view(batch_size, 64, 1, 1).repeat(1, 1, 8, 8)
+            decoder_input = torch.cat([h, style_expanded], dim=1)
             
             frame = self.decoder(decoder_input)
             frames.append(frame)
+            
+            # Обновляем вход
+            next_features = self.encoder(frame)
+            next_features = next_features.view(batch_size, -1)
+            lstm_input = next_features.unsqueeze(1)
         
         return torch.stack(frames, dim=1)
 
@@ -227,34 +214,34 @@ class NeuralAnimator:
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         print(f"🔥 Устройство: {self.device}")
         
-        # Загружаем модели
+        # Создаём модели
         self.animator = FullNeuralAnimator().to(self.device)
         self.styler = StyleTransferAnimator().to(self.device)
         
-        # Пробуем загрузить обученные модели
+        # Пробуем загрузить
         self.load_models()
         
         self.animator.eval()
         self.styler.eval()
     
     def load_models(self):
-        """Загружает или создаёт модели"""
         models_dir = 'neural_models'
         os.makedirs(models_dir, exist_ok=True)
         
-        # Если нет моделей - используем случайные (но они будут работать!)
         if os.path.exists(f'{models_dir}/animator.pth'):
             self.animator.load_state_dict(torch.load(f'{models_dir}/animator.pth', map_location=self.device))
             print("✅ Аниматор загружен")
         else:
-            print("⚠️ Модель не найдена, используется случайная (всё равно работает!)")
+            print("⚠️ Модель не найдена, используем случайную")
         
         if os.path.exists(f'{models_dir}/styler.pth'):
             self.styler.load_state_dict(torch.load(f'{models_dir}/styler.pth', map_location=self.device))
             print("✅ Стилизатор загружен")
     
-    def generate_animation(self, image, style='wave', num_frames=25, size=256):
-        """Генерирует анимацию полностью нейросетью"""
+    def generate_animation(self, image, style='wave', num_frames=20, size=128):
+        """Генерирует анимацию"""
+        if image is None:
+            return None
         
         # Подготовка
         if isinstance(image, np.ndarray):
@@ -266,7 +253,7 @@ class NeuralAnimator:
         img_tensor = torch.from_numpy(np.array(img)).float() / 127.5 - 1
         img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).to(self.device)
         
-        # === ВСЁ ДЕЛАЕТ НЕЙРОСЕТЬ ===
+        # ВСЁ ДЕЛАЕТ НЕЙРОСЕТЬ
         with torch.no_grad():
             if style in ['wave', 'pulse', 'glitch', 'melt', 'twist']:
                 frames_tensor = self.styler(img_tensor, style=style, num_frames=num_frames)
@@ -286,99 +273,75 @@ class NeuralAnimator:
         
         return temp_file.name
 
-# ============ ОБУЧЕНИЕ НЕЙРОСЕТИ ============
+# ============ ОБУЧЕНИЕ ============
 def train_neural_animator():
-    """Полностью нейросетевое обучение"""
-    print("🧠 Обучаем нейросеть делать ВСЁ...")
+    """Обучение нейросети"""
+    print("🧠 Обучаем нейросеть...")
     
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
-    # Создаём модели
     animator = FullNeuralAnimator().to(device)
     styler = StyleTransferAnimator().to(device)
     
-    # Оптимизаторы
     opt_anim = torch.optim.Adam(animator.parameters(), lr=0.0001)
     opt_style = torch.optim.Adam(styler.parameters(), lr=0.0001)
     
-    # Функция потерь
     mse = nn.MSELoss()
     
     print("🚀 Начинаем обучение...")
     
     for epoch in range(10):
-        # Генерируем случайные данные
-        batch_size = 4
-        
-        # 1. Создаём случайные изображения
+        batch_size = 2
         fake_images = torch.randn(batch_size, 3, 128, 128, device=device)
         
-        # 2. Обучаем FullNeuralAnimator
+        # Обучаем аниматор
         frames = animator(fake_images, num_frames=15)
-        
-        # Потери: плавность + согласованность
-        loss_smooth = mse(frames[:, 1:], frames[:, :-1])  # Соседние кадры похожи
-        loss_consistency = mse(frames.mean(dim=1), fake_images)  # Среднее похоже на оригинал
+        loss_smooth = mse(frames[:, 1:], frames[:, :-1])
+        loss_consistency = mse(frames.mean(dim=1), fake_images)
         loss_anim = loss_smooth + 0.5 * loss_consistency
         
         opt_anim.zero_grad()
         loss_anim.backward()
         opt_anim.step()
         
-        # 3. Обучаем StyleTransferAnimator
-        for style in ['wave', 'pulse', 'glitch', 'melt', 'twist']:
-            style_frames = styler(fake_images, style=style, num_frames=15)
-            
-            loss_style_smooth = mse(style_frames[:, 1:], style_frames[:, :-1])
-            loss_style_consistency = mse(style_frames.mean(dim=1), fake_images)
-            loss_style = loss_style_smooth + 0.5 * loss_style_consistency
-            
-            opt_style.zero_grad()
-            loss_style.backward()
-            opt_style.step()
+        # Обучаем стилизатор
+        style = 'wave'
+        style_frames = styler(fake_images, style=style, num_frames=15)
+        loss_style_smooth = mse(style_frames[:, 1:], style_frames[:, :-1])
+        loss_style_consistency = mse(style_frames.mean(dim=1), fake_images)
+        loss_style = loss_style_smooth + 0.5 * loss_style_consistency
+        
+        opt_style.zero_grad()
+        loss_style.backward()
+        opt_style.step()
         
         print(f"Epoch {epoch+1}/10 | Loss: {loss_anim.item():.4f} | Style: {loss_style.item():.4f}")
     
-    # Сохраняем модели
     os.makedirs('neural_models', exist_ok=True)
     torch.save(animator.state_dict(), 'neural_models/animator.pth')
     torch.save(styler.state_dict(), 'neural_models/styler.pth')
     
-    print("✅ Обучение завершено! Модели сохранены.")
-    return animator, styler
+    print("✅ Обучение завершено!")
+    return "✅ Модель обучена!"
 
 # ============ GRADIO ИНТЕРФЕЙС ============
 animator = NeuralAnimator()
 
-def generate_distortion(image, style, frames, size):
-    """Функция для Gradio"""
+def generate_wrapper(image, style, frames, size):
     if image is None:
         return None
-    
     try:
-        gif_path = animator.generate_animation(
-            image,
-            style=style,
-            num_frames=int(frames),
-            size=int(size)
-        )
-        return gif_path
+        return animator.generate_animation(image, style, int(frames), int(size))
     except Exception as e:
         print(f"Ошибка: {e}")
         return None
 
 # Создаём интерфейс
-with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая анимация") as demo:
+with gr.Blocks(title="🧠 Нейросетевая анимация") as demo:
     gr.Markdown("""
     # 🧠 ПОЛНОСТЬЮ НЕЙРОСЕТЕВАЯ АНИМАЦИЯ
     
-    ### Нейросеть делает ВСЁ:
-    - 🎨 Анализирует структуру изображения
-    - 🧮 Предсказывает движение
-    - 🎬 Генерирует каждый кадр
-    - ⏱️ Создаёт временную последовательность
-    
-    **Никаких ручных алгоритмов — только нейросеть!**
+    Нейросеть делает ВСЁ: анализ, предсказание движения, генерацию кадров!
     """)
     
     with gr.Row():
@@ -386,7 +349,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая ан
             input_image = gr.Image(
                 label="📸 Загрузи фото",
                 type="numpy",
-                height=400
+                height=300
             )
             
             style = gr.Dropdown(
@@ -396,7 +359,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая ан
                     ("Глитч 📺", "glitch"),
                     ("Плавление 🕯️", "melt"),
                     ("Скручивание 🌀", "twist"),
-                    ("Сюрреализм 🎭", "dream"),
                     ("Нейросетевой 🧠", "neural")
                 ],
                 label="🎨 Стиль анимации",
@@ -405,28 +367,31 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая ан
             
             frames = gr.Slider(
                 minimum=10,
-                maximum=40,
+                maximum=30,
                 value=20,
                 step=5,
                 label="Количество кадров"
             )
             
             size = gr.Slider(
-                minimum=128,
-                maximum=512,
-                value=256,
+                minimum=64,
+                maximum=256,
+                value=128,
                 step=64,
-                label="Размер (качество/скорость)"
+                label="Размер (чем меньше, тем быстрее)"
             )
             
-            generate_btn = gr.Button("🧠 Запустить нейросеть!", variant="primary", size="lg")
-            train_btn = gr.Button("🎓 Обучить нейросеть", variant="secondary", size="sm")
+            with gr.Row():
+                generate_btn = gr.Button("🧠 Запустить!", variant="primary")
+                train_btn = gr.Button("🎓 Обучить", variant="secondary")
+            
+            status = gr.Textbox(label="Статус", value="Готов к работе")
         
         with gr.Column(scale=1):
             output_gif = gr.Image(
-                label="🎬 Нейросеть сгенерировала!",
+                label="🎬 Результат",
                 type="filepath",
-                height=500
+                height=400
             )
             
             download_btn = gr.DownloadButton(
@@ -436,7 +401,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая ан
     
     # Логика
     generate_btn.click(
-        fn=generate_distortion,
+        fn=generate_wrapper,
         inputs=[input_image, style, frames, size],
         outputs=[output_gif]
     ).then(
@@ -448,34 +413,13 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧠 Нейросетевая ан
     train_btn.click(
         fn=train_neural_animator,
         inputs=[],
-        outputs=[]
-    ).then(
-        fn=lambda: "✅ Модель обучена! Перезапустите анимацию.",
-        inputs=[],
-        outputs=[gr.Textbox(label="Статус")]
+        outputs=[status]
     )
-    
-    gr.Markdown("""
-    ### 🔬 Как это работает
-    
-    1. **Нейросеть-кодировщик** понимает структуру изображения
-    2. **LSTM-слой** запоминает как меняется анимация во времени
-    3. **Нейросеть-декодер** генерирует каждый кадр
-    4. **Векторы стиля** управляют типом анимации
-    
-    **ВСЁ ОБУЧАЕТСЯ НЕЙРОСЕТЬЮ!**
-    """)
 
 if __name__ == "__main__":
     print("""
-    🧠 ЗАПУСКАЕМ ПОЛНОСТЬЮ НЕЙРОСЕТЕВУЮ АНИМАЦИЮ!
+    🧠 ЗАПУСКАЕМ НЕЙРОСЕТЕВУЮ АНИМАЦИЮ!
     📱 Открой браузер: http://localhost:7860
-    
-    Нейросеть делает ВСЁ:
-    - Анализ фото
-    - Предсказание движения
-    - Генерация кадров
-    - Создание анимации
     """)
     
     demo.launch(