from transformers import pipeline
from vista3d_config import VISTA3DConfig
from vista3d_model import VISTA3DModel, register_my_model
from vista3d_pipeline import VISTA3DPipeline, register_simple_pipeline


class HuggingFacePipelineHelper:

    def __init__(self, pipeline_name: str = "vista3d"):
        self.pipeline_name = pipeline_name

    def __model_register(self):
        register_my_model()

    def __pipeline_register(self):
        register_simple_pipeline()

    def get_pipeline(self):
        self.__model_register()
        self.__pipeline_register()
        return pipeline(self.pipeline_name)

    def _update_config(self, config, config_dict):
        if config_dict:
            for key in config_dict:
                if hasattr(config, key) and getattr(config, key) != config_dict[key]:
                    setattr(config, key, config_dict[key])
        return config

    def init_pipeline(self, pretrained_model_name_or_path: str, **kwargs):
        config = VISTA3DConfig()
        config_dict = kwargs.pop("config_dict", None)
        self._update_config(config, config_dict)
        model = VISTA3DModel(config)
        model.from_pretrained(
            pretrained_model_name_or_path=pretrained_model_name_or_path
        )
        return VISTA3DPipeline(model, **kwargs)