Spaces:

sitammeur
/

TextSnap

Running on Zero

Update src/florence/model.py

40cf531 verified 4 months ago

3.84 kB

# Importing necessary libraries
import sys
from typing import Optional
from unittest.mock import patch
from PIL import Image
import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from transformers.dynamic_module_utils import get_imports
import spaces
import gradio as gr

# Local imports
from src.logger import logging
from src.exception import CustomExceptionHandling


# Florence-2's custom modeling file (modeling_florence2.py) unconditionally
# lists "flash_attn" as a required import, even though we load the model
# with attn_implementation="eager" (which never touches flash_attn).
# Actually installing a working flash-attn build inside a Space is slow
# and often fails outright, since it needs a prebuilt wheel that exactly
# matches the image's torch/CUDA/Python ABI -- if no match exists, pip
# silently produces a non-functional stub package, which still leaves
# the model unable to load. Since flash_attn isn't actually needed here,
# the standard fix is to patch transformers' import-checker to drop
# "flash_attn" from the required-imports list, only for this one file.
# Reference: https://huggingface.co/microsoft/Florence-2-large-ft/discussions/4
def _fixed_get_imports(filename):
    imports = get_imports(filename)
    if str(filename).endswith("/modeling_florence2.py") and "flash_attn" in imports:
        imports.remove("flash_attn")
    return imports


# Load model and processor from Hugging Face
model_id = "microsoft/Florence-2-large-ft"
try:
    with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports):
        model = (
            AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager", torch_dtype=torch.float16, trust_remote_code=True)
            .to("cuda")
            .eval()
        )
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    logging.info("Model and processor loaded successfully.")

# Handle exceptions that may occur during the process
except Exception as e:
    # Custom exception handling
    raise CustomExceptionHandling(e, sys) from e


@spaces.GPU
def run_example(
    task_prompt: str, image: Image.Image, text_input: Optional[str] = None
) -> str:
    """
    Runs an example using the given task prompt and image.
    Args:
        - task_prompt (str): The task prompt for the example.
        - image (PIL.Image.Image): The image to be processed.
        - text_input (str, optional): Additional text input to be appended to the task prompt. Defaults to None.
    Returns:
        str: The parsed answer generated by the model.
    """
    try:
        # Check if image is None
        if image is None:
            gr.Warning("Please provide an image.")

        # If there is no text input, use the task prompt as the prompt
        prompt = task_prompt if text_input is None else task_prompt + text_input

        # Process the image and text input
        inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.float16)

        # Generate the answer using the model
        generated_ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )
        generated_text = processor.batch_decode(
            generated_ids, skip_special_tokens=False
        )[0]
        parsed_answer = processor.post_process_generation(
            generated_text, task=task_prompt, image_size=(image.width, image.height)
        )

        # Return the parsed answer
        return parsed_answer

    # Handle exceptions that may occur during the process
    except Exception as e:
        # Custom exception handling
        raise CustomExceptionHandling(e, sys) from e

	# Importing necessary libraries
	import sys
	from typing import Optional
	from unittest.mock import patch
	from PIL import Image
	import torch
	from transformers import AutoProcessor, AutoModelForCausalLM
	from transformers.dynamic_module_utils import get_imports
	import spaces
	import gradio as gr

	# Local imports
	from src.logger import logging
	from src.exception import CustomExceptionHandling


	# Florence-2's custom modeling file (modeling_florence2.py) unconditionally
	# lists "flash_attn" as a required import, even though we load the model
	# with attn_implementation="eager" (which never touches flash_attn).
	# Actually installing a working flash-attn build inside a Space is slow
	# and often fails outright, since it needs a prebuilt wheel that exactly
	# matches the image's torch/CUDA/Python ABI -- if no match exists, pip
	# silently produces a non-functional stub package, which still leaves
	# the model unable to load. Since flash_attn isn't actually needed here,
	# the standard fix is to patch transformers' import-checker to drop
	# "flash_attn" from the required-imports list, only for this one file.
	# Reference: https://huggingface.co/microsoft/Florence-2-large-ft/discussions/4
	def _fixed_get_imports(filename):
	imports = get_imports(filename)
	if str(filename).endswith("/modeling_florence2.py") and "flash_attn" in imports:
	imports.remove("flash_attn")
	return imports


	# Load model and processor from Hugging Face
	model_id = "microsoft/Florence-2-large-ft"
	try:
	with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports):
	model = (
	AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager", torch_dtype=torch.float16, trust_remote_code=True)
	.to("cuda")
	.eval()
	)
	processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
	logging.info("Model and processor loaded successfully.")

	# Handle exceptions that may occur during the process
	except Exception as e:
	# Custom exception handling
	raise CustomExceptionHandling(e, sys) from e


	@spaces.GPU
	def run_example(
	task_prompt: str, image: Image.Image, text_input: Optional[str] = None
	) -> str:
	"""
	Runs an example using the given task prompt and image.

	Args:
	- task_prompt (str): The task prompt for the example.
	- image (PIL.Image.Image): The image to be processed.
	- text_input (str, optional): Additional text input to be appended to the task prompt. Defaults to None.

	Returns:
	str: The parsed answer generated by the model.
	"""
	try:
	# Check if image is None
	if image is None:
	gr.Warning("Please provide an image.")

	# If there is no text input, use the task prompt as the prompt
	prompt = task_prompt if text_input is None else task_prompt + text_input

	# Process the image and text input
	inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.float16)

	# Generate the answer using the model
	generated_ids = model.generate(
	input_ids=inputs["input_ids"],
	pixel_values=inputs["pixel_values"],
	max_new_tokens=1024,
	early_stopping=False,
	do_sample=False,
	num_beams=3,
	)
	generated_text = processor.batch_decode(
	generated_ids, skip_special_tokens=False
	)[0]
	parsed_answer = processor.post_process_generation(
	generated_text, task=task_prompt, image_size=(image.width, image.height)
	)

	# Return the parsed answer
	return parsed_answer

	# Handle exceptions that may occur during the process
	except Exception as e:
	# Custom exception handling
	raise CustomExceptionHandling(e, sys) from e