38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
|
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||
|
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import io
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def load_tokenizer(name_or_path: str, **kwargs):
|
||
|
|
return AutoTokenizer.from_pretrained(name_or_path, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def load_processor(name_or_path: str, **kwargs):
|
||
|
|
try:
|
||
|
|
proc = AutoProcessor.from_pretrained(name_or_path, **kwargs)
|
||
|
|
except (OSError, ValueError) as e:
|
||
|
|
logger.warning(f"Failed to load processor from {name_or_path}: {e}")
|
||
|
|
proc = None
|
||
|
|
|
||
|
|
# If HF returned a tokenizer, discard it.
|
||
|
|
if isinstance(proc, PreTrainedTokenizerBase) or not isinstance(proc, ProcessorMixin):
|
||
|
|
proc = None
|
||
|
|
|
||
|
|
return proc
|
||
|
|
|
||
|
|
|
||
|
|
def encode_image_for_rollout_engine(image) -> str:
|
||
|
|
"""Load an image from path, ensure RGB, encode as PNG base64 string."""
|
||
|
|
buffer = io.BytesIO()
|
||
|
|
if image.mode != "RGB":
|
||
|
|
image = image.convert("RGB")
|
||
|
|
image.save(buffer, format="PNG")
|
||
|
|
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|