Update app.py

app.py

@@ -7,101 +7,7 @@ from huggingface_hub import hf_hub_download
 import torch.nn as nn
 
 class SpriteGenerator(nn.Module):
-    def __init__(self, text_encoder_name="t5-base", latent_dim=512):
-        super(SpriteGenerator, self).__init__()
-        
-        # 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
-            
-        # Proiezione dal testo al latent space
-        self.text_projection = nn.Sequential(
-            nn.Linear(768, latent_dim),
-            nn.LeakyReLU(0.2),
-            nn.Linear(latent_dim, latent_dim)
-        )
-        
-        # Generator
-        self.generator = nn.Sequential(
-            # Input: latent_dim x 1 x 1 -> 512 x 4 x 4
-            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(512),
-            nn.ReLU(True),
-
-            # 512 x 4 x 4 -> 256 x 8 x 8
-            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(256),
-            nn.ReLU(True),
-
-            # 256 x 8 x 8 -> 128 x 16 x 16
-            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(128),
-            nn.ReLU(True),
-
-            # 128 x 16 x 16 -> 64 x 32 x 32
-            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(64),
-            nn.ReLU(True),
-
-            # 64 x 32 x 32 -> 32 x 64 x 64
-            nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(32),
-            nn.ReLU(True),
-
-            # 32 x 64 x 64 -> 16 x 128 x 128
-            nn.ConvTranspose2d(32, 16, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(16),
-            nn.ReLU(True),
-
-            # 16 x 128 x 128 -> 3 x 256 x 256
-            nn.ConvTranspose2d(16, 3, 4, 2, 1, bias=False),
-        )
-        
-        # Frame interpolator
-        self.frame_interpolator = nn.Sequential(
-            nn.Linear(latent_dim + 1, latent_dim),
-            nn.LeakyReLU(0.2),
-            nn.Linear(latent_dim, latent_dim),
-            nn.LeakyReLU(0.2)
-        )
-
-    def forward(self, input_ids, attention_mask, num_frames=1):
-        batch_size = input_ids.shape[0]
-        
-        # Encode text usando il T5 completo
-        text_outputs = self.text_encoder.encoder(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            return_dict=True
-        )
-        
-        # Get text features
-        text_features = text_outputs.last_hidden_state.mean(dim=1)
-        
-        # Project to latent space
-        latent_vector = self.text_projection(text_features)
-        
-        # Generate multiple frames if needed
-        all_frames = []
-        for frame_idx in range(max(num_frames.max().item(), 1)):
-            frame_info = torch.ones((batch_size, 1), device=latent_vector.device) * frame_idx / max(num_frames.max().item(), 1)
-            
-            # Combine latent vector with frame info
-            frame_latent = self.frame_interpolator(
-                torch.cat([latent_vector, frame_info], dim=1)
-            )
-            
-            # Generate frame
-            frame_latent_reshaped = frame_latent.unsqueeze(2).unsqueeze(3)
-            frame = self.generator(frame_latent_reshaped)
-            frame = torch.tanh(frame)
-            all_frames.append(frame)
-        
-        # Stack all frames
-        sprites = torch.stack(all_frames, dim=1)
-        
-        return sprites
+    # ... (la classe SpriteGenerator rimane invariata) ...
 
 def initialize_model():
     print("Inizializzazione del modello...")
@@ -110,12 +16,19 @@ def initialize_model():
     model = SpriteGenerator()
     
     try:
+        # Scarica il modello da Hugging Face Hub
+        model_path = hf_hub_download(
+            repo_id="Lod34/Animator2D-v2",
+            filename="pytorch_model.bin",
+            repo_type="model"
+        )
+        
         # Carica il modello
-        state_dict = torch.load("Animator2D-v2.pth", map_location=device)
+        state_dict = torch.load(model_path, map_location=device)
         model.load_state_dict(state_dict)
         model = model.to(device)
         model.eval()
-        print("Modello caricato con successo!")
+        print("Modello caricato con successo da Hugging Face Hub!")
         return model, device
     except Exception as e:
         print(f"Errore nel caricamento del modello: {str(e)}")