Repository: vllm-project/vllm
File: tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py (Python)

Code before:
            if match:
                year = int(match.group(1)[:4])
                julian_day = match.group(1).split("T")[0][4:]
                if len(julian_day) == 3:
                    julian_day = int(julian_day)
                else:
                    julian_day = (
                        datetime.datetime.strptime(julian_day, "%m%d")
                        .timetuple()
                        .tm_yday
                    )
                temporal_coords.append([year, julian_day])
        except Exception:
            logger.exception("Could not extract timestamp for %s", file)

    imgs = np.stack(imgs, axis=0)  # num_frames, H, W, C
    imgs = np.moveaxis(imgs, -1, 0).astype("float32")  # C, num_frames, H, W
    imgs = np.expand_dims(imgs, axis=0)  # add batch di

    return imgs, temporal_coords, location_coords, metas


class PrithviMultimodalDataProcessor(IOProcessor[ImagePrompt, ImageRequestOutput]):
    indices = [0, 1, 2, 3, 4, 5]

    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer | None = None):
        super().__init__(vllm_config)

        self.datamodule = Sen1Floods11NonGeoDataModule(
            data_root=datamodule_config["data_root"],
            batch_size=datamodule_config["batch_size"],
            num_workers=datamodule_config["num_workers"],
            bands=datamodule_config["bands"],
            drop_last=datamodule_config["drop_last"],
            test_transform=datamodule_config["test_transform"],
        )
        self.img_size = 512
        self.h1 = 1
        self.w1 = 1
        self.original_h = 512
        self.original_w = 512
        self.batch_size = 1
        self.meta_data = None
        self.requests_cache: dict[str, dict[str, Any]] = {}
        self.indices = DEFAULT_INPUT_INDICES

    def parse_data(self, data: object) -> ImagePrompt:
        if isinstance(data, dict):
            return ImagePrompt(**data)

        raise ValueError("Prompt data should be an `ImagePrompt`")

Review comment (suggestion):
Can the renderer be `None`?. From @DarkLight1337's comment I understand all models must have a renderer.

Code after:
            if match:
                year = int(match.group(1)[:4])
                julian_day = match.group(1).split("T")[0][4:]
                if len(julian_day) == 3:
                    julian_day = int(julian_day)
                else:
                    julian_day = (
                        datetime.datetime.strptime(julian_day, "%m%d")
                        .timetuple()
                        .tm_yday
                    )
                temporal_coords.append([year, julian_day])
        except Exception:
            logger.exception("Could not extract timestamp for %s", file)

    imgs = np.stack(imgs, axis=0)  # num_frames, H, W, C
    imgs = np.moveaxis(imgs, -1, 0).astype("float32")  # C, num_frames, H, W
    imgs = np.expand_dims(imgs, axis=0)  # add batch di

    return imgs, temporal_coords, location_coords, metas


class PrithviMultimodalDataProcessor(IOProcessor[ImagePrompt, ImageRequestOutput]):
    indices = [0, 1, 2, 3, 4, 5]

    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
        super().__init__(vllm_config, renderer)

        self.datamodule = Sen1Floods11NonGeoDataModule(
            data_root=datamodule_config["data_root"],
            batch_size=datamodule_config["batch_size"],
            num_workers=datamodule_config["num_workers"],
            bands=datamodule_config["bands"],
            drop_last=datamodule_config["drop_last"],
            test_transform=datamodule_config["test_transform"],
        )
        self.img_size = 512
        self.h1 = 1
        self.w1 = 1
        self.original_h = 512
        self.original_w = 512
        self.batch_size = 1
        self.meta_data = None
        self.requests_cache: dict[str, dict[str, Any]] = {}
        self.indices = DEFAULT_INPUT_INDICES

    def parse_data(self, data: object) -> ImagePrompt:
        if isinstance(data, dict):
            return ImagePrompt(**data)

        raise ValueError("Prompt data should be an `ImagePrompt`")