Errors with Attention Mechanisms in Janus Inference Script (FlashAttention 2.0 & _flash_supports_window_size)

[AlanPonnachan]

I encountered multiple issues when running the Janus-1.3B model inference script, both with and without enabling FlashAttention. These errors prevent successful execution of the model in a standard environment, such as Google Colab. Below are the details of the issues:

Problem 1: Error Without FlashAttention

When running the inference script without FlashAttention, the following error is raised during the generate function call:
NameError: name '_flash_supports_window_size' is not defined

Problem 2: FlashAttention 2.0 Unsupported

To address the above error, I attempted to install FlashAttention-2 and enable it via the attn_implementation="flash_attention_2" argument. However, this raises the following error:
ValueError: MultiModalityCausalLM does not support Flash Attention 2.0 yet. Please request to add support where the model is hosted, on its model hub page: https://huggingface.co/deepseek-ai/Janus-1.3B/discussions/new

Steps to Reproduce:

1. Install FlashAttention 2.0 via:
   pip install flash-attn --no-build-isolation
   2.Modify the script to include:

vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True, attn_implementation="flash_attention_2"
)

Thank you for your assistance! Let me know if you require additional details or logs to reproduce the issues.