Update app.py

app.py

@@ -2,17 +2,16 @@ import gradio as gr
 import torch
 from PIL import Image
 import os
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, T5ForConditionalGeneration
 from huggingface_hub import hf_hub_download
 import torch.nn as nn
 
-# Definizione del modello
 class SpriteGenerator(nn.Module):
     def __init__(self, text_encoder_name="t5-base", latent_dim=512):
         super(SpriteGenerator, self).__init__()
         
-        # Text encoder (T5)
-        self.text_encoder = AutoModel.from_pretrained(text_encoder_name)
+        # Text encoder (T5 with lm_head)
+        self.text_encoder = T5ForConditionalGeneration.from_pretrained(text_encoder_name)
         for param in self.text_encoder.parameters():
             param.requires_grad = False
             
@@ -23,30 +22,39 @@ class SpriteGenerator(nn.Module):
             nn.Linear(latent_dim, latent_dim)
         )
         
-        # Generator
+        # Generator modificato per corrispondere ai pesi salvati
         self.generator = nn.Sequential(
-            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),
+            # Input: latent_dim x 1 x 1
+            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),  # -> 512 x 4 x 4
             nn.BatchNorm2d(512),
-            nn.LeakyReLU(0.2, inplace=True),
-            
-            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),  # -> 256 x 8 x 8
             nn.BatchNorm2d(256),
-            nn.LeakyReLU(0.2, inplace=True),
-            
-            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),  # -> 128 x 16 x 16
             nn.BatchNorm2d(128),
-            nn.LeakyReLU(0.2, inplace=True),
-            
-            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),  # -> 64 x 32 x 32
             nn.BatchNorm2d(64),
-            nn.LeakyReLU(0.2, inplace=True),
-            
-            nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),  # -> 32 x 64 x 64
             nn.BatchNorm2d(32),
-            nn.LeakyReLU(0.2, inplace=True),
+            nn.ReLU(True),
+
+            nn.ConvTranspose2d(32, 16, 4, 2, 1, bias=False),  # -> 16 x 128 x 128
+            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.ConvTranspose2d(32, 3, 4, 2, 1, bias=False),
-            nn.Tanh()
+            nn.Conv2d(16, 3, 3, 1, 1)  # Layer di output per RGB
         )
         
         # Frame interpolator
@@ -60,8 +68,8 @@ class SpriteGenerator(nn.Module):
     def forward(self, input_ids, attention_mask, num_frames=1):
         batch_size = input_ids.shape[0]
         
-        # Encode text
-        text_outputs = self.text_encoder(
+        # Encode text usando il T5 completo
+        text_outputs = self.text_encoder.encoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
             return_dict=True
@@ -92,7 +100,7 @@ class SpriteGenerator(nn.Module):
         sprites = torch.stack(all_frames, dim=1)
         
         return sprites
-
+        
 # Costanti
 MODEL_ID = "Lod34/Animator2D-v2"
 CACHE_DIR = "model_cache"
@@ -242,5 +250,3 @@ def create_interface():
 
 # Crea l'interfaccia
 demo = create_interface()
-
-# Per Spaces, non usare demo.launch()