import modal
from smolagents import Tool

from transformers import pipeline

from .app import app
from .image import image


@app.cls(gpu="T4", image=image)
class InferencePipelineModalApp:
    @modal.method()
    def forward(self, task: str, model_name: str, image, **kwargs):
        try:
            pipe = pipeline(task=task, model=model_name, device="cuda")
            result = pipe(image, **kwargs)
            return result
        except Exception as e:
            raise Exception(f"Error in inference pipeline: {e}. Please try again with a different model.")