Update modeling_pixel.py

modeling_pixel.py

@@ -18,22 +18,26 @@ class ResidualBlock(nn.Module):
 
 class TopAIImageGenerator(PreTrainedModel):
     config_class = TopAIImageConfig
-    # תיקון השגיאה: מגדירים כרשימה ריקה, Transformers כבר יסתדר עם זה
-    all_tied_weights_keys = [] 
+    
+    # תיקון קריטי: הגדרה כמילון ריק כדי למנוע את ה-AttributeError
+    all_tied_weights_keys = {} 
 
     def __init__(self, config):
         super().__init__(config)
-        # סינכרון המימדים לפי ה-Checkpoint שלכם
         self.text_projection = nn.Linear(config.input_dim, 4 * 4 * config.hidden_dim)
         
+        # בנייה מחדש של ה-Decoder כדי להתאים למבנה ה-Checkpoint שלכם
         self.decoder = nn.Sequential(
-            self._upsample(config.hidden_dim, 256), # 4x4 -> 8x8
-            ResidualBlock(256),
-            self._upsample(256, 128),              # 8x8 -> 16x16
-            ResidualBlock(128),
-            self._upsample(128, 64),               # 16x16 -> 32x32
-            self._upsample(64, 32),                # 32x32 -> 64x64
-            nn.Conv2d(32, config.image_channels, kernel_size=3, padding=1),
+            self._upsample(config.hidden_dim, 256), # 0
+            ResidualBlock(256),                    # 1
+            self._upsample(256, 128),              # 2
+            ResidualBlock(128),                    # 3
+            self._upsample(128, 64),               # 4
+            self._upsample(64, 32),                # 5
+            nn.Conv2d(32, 32, kernel_size=3, padding=1), # 6
+            nn.BatchNorm2d(32),                    # 7
+            nn.ReLU(True),                         # 8
+            nn.Conv2d(32, config.image_channels, kernel_size=3, padding=1), # 9
             nn.Tanh()
         )