diff --git a/src/opengeodeweb_viewer/object/object_methods.py b/src/opengeodeweb_viewer/object/object_methods.py
index d40df03..215aaba 100644
--- a/src/opengeodeweb_viewer/object/object_methods.py
+++ b/src/opengeodeweb_viewer/object/object_methods.py
@@ -108,20 +108,28 @@ def SetPointsColor(self, data_id: str, red: int, green: int, blue: int) -> None:
     def SetBlocksVisibility(
         self, data_id: str, block_ids: list[int], visibility: bool
     ) -> None:
-        mapper = self.get_vtk_pipeline(data_id).mapper
+        pipeline = self.get_vtk_pipeline(data_id)
+        mapper = pipeline.mapper
         if not isinstance(mapper, vtkCompositePolyDataMapper):
             raise Exception("Mapper is not a vtkCompositePolyDataMapper")
+        blocks = pipeline.blockDataSets
+        attributes = mapper.GetCompositeDataDisplayAttributes()
         for block_id in block_ids:
-            mapper.SetBlockVisibility(block_id, visibility)
+            attributes.SetBlockVisibility(blocks[block_id], visibility)
 
     def SetBlocksColor(
         self, data_id: str, block_ids: list[int], red: int, green: int, blue: int
     ) -> None:
-        mapper = self.get_vtk_pipeline(data_id).mapper
+        pipeline = self.get_vtk_pipeline(data_id)
+        mapper = pipeline.mapper
         if not isinstance(mapper, vtkCompositePolyDataMapper):
             raise Exception("Mapper is not a vtkCompositePolyDataMapper")
+        blocks = pipeline.blockDataSets
+        attributes = mapper.GetCompositeDataDisplayAttributes()
         for block_id in block_ids:
-            mapper.SetBlockColor(block_id, [red / 255, green / 255, blue / 255])
+            attributes.SetBlockColor(
+                blocks[block_id], [red / 255, green / 255, blue / 255]
+            )
 
     def clearColors(self, data_id: str) -> None:
         db = self.get_vtk_pipeline(data_id)
diff --git a/src/opengeodeweb_viewer/rpc/model/model_protocols.py b/src/opengeodeweb_viewer/rpc/model/model_protocols.py
index 6ac4705..6ca4794 100644
--- a/src/opengeodeweb_viewer/rpc/model/model_protocols.py
+++ b/src/opengeodeweb_viewer/rpc/model/model_protocols.py
@@ -51,6 +51,16 @@ def registerModel(self, rpc_params: RpcParams) -> None:
             attributes = vtkCompositeDataDisplayAttributes()
             mapper.SetCompositeDataDisplayAttributes(attributes)
             data = VtkPipeline(reader, mapper, filter)
+            iterator = geometry_output.NewTreeIterator()
+            iterator.InitTraversal()
+            while not iterator.IsDoneWithTraversal():
+                block = iterator.GetCurrentDataObject()
+                if block:
+                    flat_index = iterator.GetCurrentFlatIndex()
+                    while flat_index > len(data.blockDataSets):
+                        data.blockDataSets.append(None)
+                    data.blockDataSets.append(block)
+                iterator.GoToNextItem()
             self.registerObject(data_id, file_name, data)
         except Exception as e:
             print(f"Error registering model {data_id}: {str(e)}", flush=True)
diff --git a/src/opengeodeweb_viewer/vtk_protocol.py b/src/opengeodeweb_viewer/vtk_protocol.py
index 7968afe..4fb092a 100644
--- a/src/opengeodeweb_viewer/vtk_protocol.py
+++ b/src/opengeodeweb_viewer/vtk_protocol.py
@@ -18,6 +18,7 @@
     vtkRenderWindow,
     vtkCompositePolyDataMapper,
 )
+from vtkmodules.vtkCommonDataModel import vtkDataObject
 from vtkmodules.vtkRenderingAnnotation import vtkCubeAxesActor, vtkAxesActor
 from vtkmodules.vtkInteractionWidgets import vtkOrientationMarkerWidget
 
@@ -41,6 +42,7 @@ class VtkPipeline:
     mapper: vtkMapper
     filter: vtkAlgorithm | None = None
     actor: vtkActor = field(default_factory=vtkActor)
+    blockDataSets: list[vtkDataObject | None] = field(default_factory=list)
 
 
 class VtkTypingMixin: