Large scale controlnet

components/prompt_based_laion_retrieval/fondant_component.yaml

@@ -25,3 +25,7 @@ args:
   aesthetic_weight:
     description: Weight of the aesthetic embedding when added to the query, between 0 and 1
     type: float
+  url:
+    description: The url of the backend clip retrieval service, defaults to the public service
+    type: str
+    default: https://knn.laion.ai/knn-service

components/prompt_based_laion_retrieval/src/main.py

@@ -21,6 +21,7 @@ def setup(
             num_images: int,
             aesthetic_score: int,
             aesthetic_weight: float,
+            url: str,
     ) -> None:
         """
 
@@ -30,10 +31,11 @@ def setup(
                 between 0 and 9.
             aesthetic_weight: weight of the aesthetic embedding to add to the query,
                 between 0 and 1.
+            url: The url of the backend clip retrieval service, defaults to the public clip url.
         """
         self.client = ClipClient(
-            url="https://knn.laion.ai/knn-service",
-            indice_name="laion5B-L-14",
+            url=url,
+            indice_name="laion5B",  #TODO:revert back to laion5b-L-14 after backend correction
             num_images=num_images,
             aesthetic_score=aesthetic_score,
             aesthetic_weight=aesthetic_weight,

components/segment_images/src/main.py

@@ -15,15 +15,15 @@
 logger = logging.getLogger(__name__)
 
 
-def convert_to_rgb(seg: np.array):
+def convert_to_rgb(seg: np.array) -> bytes:
     """
     Converts a 2D segmentation to a RGB one which makes it possible to visualize it.
 
     Args:
         seg: 2D segmentation map as a NumPy array.
 
     Returns:
-        color_seg: 3D segmentation map contain RGB values for each pixel.
+        color_seg: the RGB segmentation map as a binary string
     """
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8,
@@ -32,9 +32,13 @@ def convert_to_rgb(seg: np.array):
     for label, color in enumerate(palette):
         color_seg[seg == label, :] = color
 
-    color_seg = color_seg.astype(np.uint8).tobytes()
+    color_seg = color_seg.astype(np.uint8)
+    image = Image.fromarray(color_seg).convert('RGB')
 
-    return color_seg
+    crop_bytes = io.BytesIO()
+    image.save(crop_bytes, format="JPEG")
+
+    return crop_bytes.getvalue()
 
 
 def process_image(image: bytes, *, processor: SegformerImageProcessor, device: str) -> torch.Tensor:
@@ -46,6 +50,7 @@ def process_image(image: bytes, *, processor: SegformerImageProcessor, device: s
         processor: The processor object for transforming the image.
         device: The device to move the transformed image to.
     """
+
     def load(img: bytes) -> Image:
         """Load the bytestring as an image."""
         bytes_ = io.BytesIO(img)

components/write_to_hf_hub/src/main.py

@@ -8,6 +8,7 @@
 
 # Define the schema for the struct using PyArrow
 import huggingface_hub
+from datasets.features.features import generate_from_arrow_type
 from PIL import Image
 
 from fondant.component import WriteComponent
@@ -71,7 +72,7 @@ def write(
                 if image_column_names and column_name in image_column_names:
                     schema_dict[column_name] = datasets.Image()
                 else:
-                    schema_dict[column_name] = datasets.Value(str(field.type.value))
+                    schema_dict[column_name] = generate_from_arrow_type(field.type.value)
 
         schema = datasets.Features(schema_dict).arrow_schema
         dataframe = dataframe[write_columns]

examples/pipelines/controlnet-interior-design/components/write_to_hub_controlnet/fondant_component.yaml

@@ -16,9 +16,7 @@ consumes:
   segmentations:
     fields:
       data:
-        type: array
-        items:
-          type: binary
+        type: binary
 
 args:
   hf_token:

examples/pipelines/controlnet-interior-design/pipeline.py

@@ -20,7 +20,12 @@
 )
 laion_retrieval_op = ComponentOp.from_registry(
     name="prompt_based_laion_retrieval",
-    arguments={"num_images": 2, "aesthetic_score": 9, "aesthetic_weight": 0.5},
+    arguments={
+        "num_images": 2,
+        "aesthetic_score": 9,
+        "aesthetic_weight": 0.5,
+        "url": None,
+    },
 )
 download_images_op = ComponentOp.from_registry(
     name="download_images",
@@ -63,8 +68,6 @@
         "hf_token": "hf_token",
         "image_column_names": ["images_data"],
     },
-    number_of_gpus=1,
-    node_pool_name="model-inference-pool",
 )
 
 pipeline = Pipeline(pipeline_name=pipeline_name, base_path=PipelineConfigs.BASE_PATH)