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 PixelConfig
+from .configuration_pixel import TopAIImageConfig # תיקון השם כאן
 
 class ResidualBlock(nn.Module):
     def __init__(self, channels):
@@ -16,21 +16,21 @@ class ResidualBlock(nn.Module):
     def forward(self, x):
         return x + self.block(x)
 
-class PixelGenerator(PreTrainedModel):
-    config_class = PixelConfig
+class TopAIImageGenerator(PreTrainedModel):
+    config_class = TopAIImageConfig
 
     def __init__(self, config):
         super().__init__(config)
         self.text_projection = nn.Linear(config.input_dim, 4 * 4 * 1024)
         
         self.decoder = nn.Sequential(
-            self._upsample(1024, 512), # 4x4 -> 8x8
+            self._upsample(1024, 512), 
             ResidualBlock(512),
-            self._upsample(512, 256),  # 8x8 -> 16x16
+            self._upsample(512, 256),  
             ResidualBlock(256),
-            self._upsample(256, 128),  # 16x16 -> 32x32
-            self._upsample(128, 64),   # 32x32 -> 64x64
-            self._upsample(64, 32),    # 64x64 -> 128x128
+            self._upsample(256, 128),  
+            self._upsample(128, 64),   
+            self._upsample(64, 32),    
             nn.Conv2d(32, config.image_channels, kernel_size=3, padding=1),
             nn.Tanh()
         )