Update app.py

app.py

@@ -8,31 +8,28 @@ import imageio
 import os
 import tempfile
 
-# ============ ИСПРАВЛЕННАЯ ПОЛНОСТЬЮ НЕЙРОСЕТЕВАЯ АРХИТЕКТУРА ============
+# ============ ИСПРАВЛЕННАЯ ВЕРСИЯ ============
 class FullNeuralAnimator(nn.Module):
-    """
-    Одна нейросеть делает ВСЁ
-    """
     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 * 8 * 8,  # 256 каналов * 8x8 (после 3х пулингов)
+            input_size=256 * 8 * 8,
             hidden_size=512,
             num_layers=2,
             batch_first=True,
             dropout=0.2
         )
         
-        # === Декодер ===
+        # Decoder
         self.dec4 = self._block(512, 256)
         self.dec3 = self._block(256, 128)
         self.dec2 = self._block(128, 64)
@@ -43,14 +40,13 @@ class FullNeuralAnimator(nn.Module):
         self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
         self.up1 = nn.ConvTranspose2d(64, 32, 2, stride=2)
         
-        # === Выход для каждого кадра ===
         self.frame_generator = nn.Sequential(
             nn.Conv2d(32, 16, 3, padding=1),
             nn.ReLU(),
             nn.Conv2d(16, 3, 3, padding=1),
             nn.Tanh()
         )
-        
+    
     def _block(self, in_ch, out_ch):
         return nn.Sequential(
             nn.Conv2d(in_ch, out_ch, 3, padding=1),
@@ -64,59 +60,52 @@ class FullNeuralAnimator(nn.Module):
     def forward(self, x, num_frames=20):
         batch_size = x.size(0)
         
-        # === 1. Кодируем изображение ===
+        # Encoder
         e1 = self.enc1(x)
         e2 = self.enc2(self.pool(e1))
         e3 = self.enc3(self.pool(e2))
         e4 = self.enc4(self.pool(e3))
         
-        # Сохраняем skip connections
         skips = [e1, e2, e3, e4]
         
-        # === 2. Подготовка для LSTM ===
-        # После 3х пулингов: 256 -> 8x8
-        bottleneck = e4.view(batch_size, -1)  # [B, 256*8*8]
+        # Bottleneck
+        bottleneck = e4.view(batch_size, -1)
         
-        # === 3. Генерируем последовательность ===
         frames = []
         hidden = None
-        
-        lstm_input = bottleneck.unsqueeze(1)  # [B, 1, features]
+        lstm_input = bottleneck.unsqueeze(1)
         
         for t in range(num_frames):
-            # LSTM предсказывает следующее состояние
             lstm_out, hidden = self.lstm(lstm_input, hidden)
             
-            # === 4. Декодируем в кадр ===
             h = lstm_out.squeeze(1).view(batch_size, 256, 8, 8)
             
-            # Декодер с skip connections
+            # Decoder с проверкой размеров
             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')
+            # Проверяем размеры и при необходимости ресайзим
+            if d4.size(2) != skips[3].size(2) or d4.size(3) != skips[3].size(3):
+                skips[3] = F.interpolate(skips[3], size=(d4.size(2), d4.size(3)), 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')
+            if d3.size(2) != skips[2].size(2) or d3.size(3) != skips[2].size(3):
+                skips[2] = F.interpolate(skips[2], size=(d3.size(2), d3.size(3)), 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')
+            if d2.size(2) != skips[1].size(2) or d2.size(3) != skips[1].size(3):
+                skips[1] = F.interpolate(skips[1], size=(d2.size(2), d2.size(3)), 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')
+            if d1.size(2) != skips[0].size(2) or d1.size(3) != skips[0].size(3):
+                skips[0] = F.interpolate(skips[0], size=(d1.size(2), d1.size(3)), mode='bilinear')
             d1 = torch.cat([d1, skips[0]], dim=1)
             d1 = self.dec1(d1)
             
-            # Генерируем кадр
             frame = self.frame_generator(d1)
             frames.append(frame)
             
@@ -132,7 +121,6 @@ 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)),
@@ -141,7 +129,6 @@ class StyleTransferAnimator(nn.Module):
             'twist': nn.Parameter(torch.randn(64)),
         })
         
-        # Encoder
         self.encoder = nn.Sequential(
             nn.Conv2d(3, 32, 4, stride=2, padding=1),
             nn.ReLU(),
@@ -153,7 +140,6 @@ class StyleTransferAnimator(nn.Module):
             nn.ReLU(),
         )
         
-        # LSTM
         self.lstm = nn.LSTM(
             input_size=256 * 8 * 8,
             hidden_size=512,
@@ -162,7 +148,6 @@ class StyleTransferAnimator(nn.Module):
             dropout=0.2
         )
         
-        # Decoder
         self.decoder = nn.Sequential(
             nn.ConvTranspose2d(512 + 64, 256, 4, stride=2, padding=1),
             nn.ReLU(),
@@ -177,11 +162,9 @@ class StyleTransferAnimator(nn.Module):
     def forward(self, x, style='wave', num_frames=20):
         batch_size = x.size(0)
         
-        # Кодируем
-        features = self.encoder(x)  # [B, 256, 8, 8]
+        features = self.encoder(x)
         features_flat = features.view(batch_size, -1)
         
-        # Стиль
         style_vector = self.style_embeddings[style]
         style_vector = style_vector.unsqueeze(0).repeat(batch_size, 1)
         
@@ -190,10 +173,8 @@ class StyleTransferAnimator(nn.Module):
         lstm_input = features_flat.unsqueeze(1)
         
         for t in range(num_frames):
-            # LSTM
             lstm_out, hidden = self.lstm(lstm_input, hidden)
             
-            # Декодируем со стилем
             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)
@@ -201,7 +182,6 @@ class StyleTransferAnimator(nn.Module):
             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)
@@ -214,11 +194,9 @@ 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()
@@ -239,11 +217,9 @@ class NeuralAnimator:
             print("✅ Стилизатор загружен")
     
     def generate_animation(self, image, style='wave', num_frames=20, size=128):
-        """Генерирует анимацию"""
         if image is None:
             return None
         
-        # Подготовка
         if isinstance(image, np.ndarray):
             img = Image.fromarray(image)
         else:
@@ -253,21 +229,18 @@ 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)
             else:
                 frames_tensor = self.animator(img_tensor, num_frames=num_frames)
         
-        # Конвертируем в кадры
         frames = []
         for t in range(num_frames):
             frame = frames_tensor[0, t].cpu().numpy().transpose(1, 2, 0)
             frame = np.clip((frame + 1) / 2, 0, 1)
             frames.append((frame * 255).astype(np.uint8))
         
-        # Сохраняем
         temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.gif')
         imageio.mimsave(temp_file.name, frames, duration=0.05, loop=0)
         
@@ -275,7 +248,6 @@ class NeuralAnimator:
 
 # ============ ОБУЧЕНИЕ ============
 def train_neural_animator():
-    """Обучение нейросети"""
     print("🧠 Обучаем нейросеть...")
     
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -290,12 +262,12 @@ def train_neural_animator():
     
     print("🚀 Начинаем обучение...")
     
-    for epoch in range(10):
+    for epoch in range(5):  # 5 эпох для быстро��о теста
         batch_size = 2
         fake_images = torch.randn(batch_size, 3, 128, 128, device=device)
         
         # Обучаем аниматор
-        frames = animator(fake_images, num_frames=15)
+        frames = animator(fake_images, num_frames=10)
         loss_smooth = mse(frames[:, 1:], frames[:, :-1])
         loss_consistency = mse(frames.mean(dim=1), fake_images)
         loss_anim = loss_smooth + 0.5 * loss_consistency
@@ -305,8 +277,7 @@ def train_neural_animator():
         opt_anim.step()
         
         # Обучаем стилизатор
-        style = 'wave'
-        style_frames = styler(fake_images, style=style, num_frames=15)
+        style_frames = styler(fake_images, style='wave', num_frames=10)
         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
@@ -315,7 +286,7 @@ def train_neural_animator():
         loss_style.backward()
         opt_style.step()
         
-        print(f"Epoch {epoch+1}/10 | Loss: {loss_anim.item():.4f} | Style: {loss_style.item():.4f}")
+        print(f"Epoch {epoch+1}/5 | 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')
@@ -399,7 +370,6 @@ with gr.Blocks(title="🧠 Нейросетевая анимация") as demo:
                 variant="primary"
             )
     
-    # Логика
     generate_btn.click(
         fn=generate_wrapper,
         inputs=[input_image, style, frames, size],