Update app.py

app.py

@@ -22,7 +22,7 @@ class SpriteGenerator(nn.Module):
             nn.Linear(latent_dim, latent_dim)
         )
         
-        # Generator modificato per corrispondere ai pesi salvati
+        # Generator modificato per corrispondere esattamente ai pesi salvati
         self.generator = nn.Sequential(
             # Input: latent_dim x 1 x 1
             nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),  # -> 512 x 4 x 4
@@ -49,12 +49,16 @@ class SpriteGenerator(nn.Module):
             nn.BatchNorm2d(16),
             nn.ReLU(True),
 
-            # Layer finale modificato per corrispondere ai pesi
             nn.ConvTranspose2d(16, 16, 4, 2, 1, bias=False),  # -> 16 x 256 x 256
             nn.BatchNorm2d(16),
             nn.ReLU(True),
             
-            nn.Conv2d(16, 3, 3, 1, 1)  # Layer di output per RGB
+            # Layer finale modificato per corrispondere esattamente ai pesi
+            nn.ConvTranspose2d(16, 3, 4, 2, 1, bias=False),  # Output layer per RGB
+            nn.BatchNorm2d(3),  # Aggiunto BatchNorm
+            nn.ReLU(True),      # Aggiunto ReLU
+            
+            nn.Conv2d(3, 3, 3, 1, 1)  # Layer di output finale
         )
         
         # Frame interpolator
@@ -100,7 +104,7 @@ class SpriteGenerator(nn.Module):
         sprites = torch.stack(all_frames, dim=1)
         
         return sprites
-        
+
 # Costanti
 MODEL_ID = "Lod34/Animator2D-v2"
 CACHE_DIR = "model_cache"