image_processor.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. from typing import cast
  2. def get_video_processor(
  3. processor_name: str,
  4. trust_remote_code: bool = False,
  5. ):
  6. """
  7. Gets a processor for the given model name via HuggingFace.
  8. """
  9. from transformers import AutoProcessor
  10. try:
  11. processor = AutoProcessor.from_pretrained(processor_name)
  12. video_processor = processor.video_processor
  13. except ValueError as e:
  14. if not trust_remote_code:
  15. err_msg = (
  16. "Failed to load the processor. If the processor is "
  17. "a custom processor not yet available in the HuggingFace "
  18. "transformers library, consider setting "
  19. "`trust_remote_code=True` in LLM or using the "
  20. "`--trust-remote-code` flag in the CLI.")
  21. raise RuntimeError(err_msg) from e
  22. else:
  23. raise e
  24. return video_processor
  25. def get_image_processor(
  26. processor_name: str,
  27. *args,
  28. trust_remote_code: bool = False,
  29. **kwargs,
  30. ):
  31. """Gets an image processor for the given model name via HuggingFace."""
  32. # don't put this import at the top level
  33. # it will call torch.cuda.device_count()
  34. from transformers import AutoImageProcessor
  35. from transformers.image_processing_utils import BaseImageProcessor
  36. try:
  37. processor = AutoImageProcessor.from_pretrained(
  38. processor_name,
  39. *args,
  40. trust_remote_code=trust_remote_code,
  41. **kwargs)
  42. except ValueError as e:
  43. # If the error pertains to the processor class not existing or not
  44. # currently being imported, suggest using the --trust-remote-code flag.
  45. # Unlike AutoTokenizer, AutoImageProcessor does not separate such errors
  46. if not trust_remote_code:
  47. err_msg = (
  48. "Failed to load the image processor. If the image processor is "
  49. "a custom processor not yet available in the HuggingFace "
  50. "transformers library, consider setting "
  51. "`trust_remote_code=True` in LLM or using the "
  52. "`--trust-remote-code` flag in the CLI.")
  53. raise RuntimeError(err_msg) from e
  54. else:
  55. raise e
  56. return cast(BaseImageProcessor, processor)