Update modeling_pixel.py

modeling_pixel.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
-from .configuration_pixel import TopAIImageConfig # תיקון השם כאן
+from .configuration_pixel import TopAIImageConfig
 
 class ResidualBlock(nn.Module):
     def __init__(self, channels):
@@ -18,19 +18,21 @@ class ResidualBlock(nn.Module):
 
 class TopAIImageGenerator(PreTrainedModel):
     config_class = TopAIImageConfig
+    all_tied_weights_keys = [] # מונע את ה-AttributeError ב-Transformers
 
     def __init__(self, config):
         super().__init__(config)
-        self.text_projection = nn.Linear(config.input_dim, 4 * 4 * 1024)
+        # תיקון המימד ל-4*4*512 בהתאם ל-Checkpoint שלכם
+        self.text_projection = nn.Linear(config.input_dim, 4 * 4 * config.hidden_dim)
         
         self.decoder = nn.Sequential(
-            self._upsample(1024, 512), 
-            ResidualBlock(512),
-            self._upsample(512, 256),  
+            self._upsample(config.hidden_dim, 256), # 4 -> 8
             ResidualBlock(256),
-            self._upsample(256, 128),  
-            self._upsample(128, 64),   
-            self._upsample(64, 32),    
+            self._upsample(256, 128),              # 8 -> 16
+            ResidualBlock(128),
+            self._upsample(128, 64),               # 16 -> 32
+            self._upsample(64, 32),                # 32 -> 64
+            self._upsample(32, 32),                # 64 -> 128 (תיקון למימד עקבי)
             nn.Conv2d(32, config.image_channels, kernel_size=3, padding=1),
             nn.Tanh()
         )
@@ -44,5 +46,5 @@ class TopAIImageGenerator(PreTrainedModel):
 
     def forward(self, text_embeddings):
         x = self.text_projection(text_embeddings)
-        x = x.view(-1, 1024, 4, 4)
+        x = x.view(-1, self.config.hidden_dim, 4, 4)
         return self.decoder(x)