from pydantic import BaseModel, Field from typing import Union, Generator, Iterator import os import logging import torch from PIL import Image # Import Image from PIL from transformers import AutoModelForCausalLM, AutoProcessor from huggingface_hub import HfApi logging.basicConfig(level=logging.INFO) class Pipe: class Valves(BaseModel): MODEL_NAME: str = Field( default="meta-llama/Llama-3.2-11B-Vision-Instruct", # Corrected model name description="Model name for Llama Vision.", ) HUGGINGFACE_API_TOKEN: str = Field( default=os.getenv("HUGGINGFACE_API_TOKEN", ""), description="API token for authenticating requests to the Hugging Face API.", ) def __init__(self): self.valves = self.Valves() self.model = None self.processor = None def validate_token(self, access_token) -> HfApi: # Replace with your actual access token api = HfApi(token=access_token) return ( api # Optionally return False if the token is invalid or an error occurred ) def load_model(self, access_token): if self.model is None or self.processor is None: # Download the model and processor self.processor = AutoProcessor.from_pretrained( self.valves.MODEL_NAME, token=access_token, # Use self.valves.MODEL_NAME ) self.model = AutoModelForCausalLM.from_pretrained( self.valves.MODEL_NAME, # Use self.valves.MODEL_NAME token=access_token, torch_dtype=torch.float16, ).cuda() def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]: access_token = self.valves.HUGGINGFACE_API_TOKEN logging.info(body) try: prompt = body.get( "prompt", "" ) # Get prompt from body or default to empty string if ( "image" in body and isinstance(body["image"], list) and len(body["image"]) > 0 ): image_path = body["image"][0] # Get the first image from the list else: image_path = None logging.info(prompt) logging.info(image_path) self.validate_token(access_token) self.load_model(access_token) logging.info("Model loaded successfully") if image_path: image = Image.open(image_path) # Open the image using PIL inputs = self.processor( images=image, text=prompt, return_tensors="pt" ).to(self.model.device) else: inputs = self.processor(text=prompt, return_tensors="pt").to( self.model.device ) output = self.model.generate(**inputs, max_new_tokens=30) return self.processor.decode(output[0], skip_special_tokens=True) except Exception as e: logging.error(f"Failed to process request: {e}") return f"Error: {e}" if __name__ == "__main__": pipe = Pipe() # Example usage result = pipe.pipe({"prompt": "Describe this image", "image": "path/to/image.jpg"}) print(result)