Add a dummy task encoder for raw image inputs

Signed-off-by: mingyuanm <[email protected]>

nemo/collections/diffusion/data/diffusion_taskencoder.py

@@ -443,3 +443,33 @@ def concat_pad(tensor_list, max_seq_length):
 
 
 pos_id_3d = PosID3D()
+
+
+
+def cook_raw_iamges(sample: dict) -> dict:
+    """
+    Processes a raw sample dictionary from energon dataset and returns a new dictionary with specific keys.
+
+    Args:
+        sample (dict): The input dictionary containing the raw sample data.
+
+    Returns:
+        dict: A new dictionary containing the processed sample data with the following keys:
+            - All keys from the result of `basic_sample_keys(sample)`
+            - 'jpg': original images
+            - 'png': contains control images
+            - 'txt': contains raw text
+    """
+    return dict(
+        **basic_sample_keys(sample),
+        images=sample['jpg'],
+        hint=sample['png'],
+        txt=sample['txt'],
+    )
+
+
+class RawImageDiffusionTaskEncoder(DefaultTaskEncoder, IOMixin):
+    cookers = [
+        # Cooker(cook),
+        Cooker(cook_raw_iamges),
+    ]