Update modeling.py

modeling.py

@@ -8,6 +8,7 @@ import torch
 from .dino_wrapper2 import DinoWrapper
 from .transformer import TriplaneTransformer
 from .synthesizer_part import TriplaneSynthesizer
+from .processor import LRMImageProcessor
 
 class CameraEmbedder(nn.Module):
     def __init__(self, raw_dim: int, embed_dim: int):
@@ -46,6 +47,8 @@ class LRMGenerator(PreTrainedModel):
     def __init__(self, config: LRMGeneratorConfig):
         super().__init__(config)
 
+        self.image_processor = LRMImageProcessor(source_size=512)
+
         self.encoder_feat_dim = config.encoder_feat_dim
         self.camera_embed_dim = config.camera_embed_dim
 
@@ -67,6 +70,10 @@ class LRMGenerator(PreTrainedModel):
         )
 
     def forward(self, image, camera):
+
+        # we use image processor directly in the forward pass
+        processed_image, source_camera = self.image_processor(image)
+        
         assert image.shape[0] == camera.shape[0], "Batch size mismatch"
         N = image.shape[0]