invoke-ai--invokeai
wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
77 行
2.9 KiB
Python
77 行
2.9 KiB
Python
from typing import Any
|
|
|
|
import torch
|
|
from PIL.Image import Image
|
|
from pydantic import field_validator
|
|
from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor
|
|
|
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
|
|
from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField, UIComponent
|
|
from invokeai.app.invocations.model import ModelIdentifierField
|
|
from invokeai.app.invocations.primitives import StringOutput
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
|
from invokeai.backend.llava_onevision_pipeline import LlavaOnevisionPipeline
|
|
from invokeai.backend.model_manager.taxonomy import ModelType
|
|
from invokeai.backend.util.devices import TorchDevice
|
|
|
|
|
|
@invocation(
|
|
"llava_onevision_vllm",
|
|
title="LLaVA OneVision VLLM",
|
|
tags=["vllm"],
|
|
category="multimodal",
|
|
version="1.0.0",
|
|
classification=Classification.Beta,
|
|
)
|
|
class LlavaOnevisionVllmInvocation(BaseInvocation):
|
|
"""Run a LLaVA OneVision VLLM model."""
|
|
|
|
images: list[ImageField] | ImageField | None = InputField(default=None, max_length=3, description="Input image.")
|
|
prompt: str = InputField(
|
|
default="",
|
|
description="Input text prompt.",
|
|
ui_component=UIComponent.Textarea,
|
|
)
|
|
vllm_model: ModelIdentifierField = InputField(
|
|
title="LLaVA Model Type",
|
|
description=FieldDescriptions.vllm_model,
|
|
ui_model_type=ModelType.LlavaOnevision,
|
|
)
|
|
|
|
@field_validator("images", mode="before")
|
|
def listify_images(cls, v: Any) -> list:
|
|
if v is None:
|
|
return v
|
|
if not isinstance(v, list):
|
|
return [v]
|
|
return v
|
|
|
|
def _get_images(self, context: InvocationContext) -> list[Image]:
|
|
if self.images is None:
|
|
return []
|
|
|
|
image_fields = self.images if isinstance(self.images, list) else [self.images]
|
|
return [context.images.get_pil(image_field.image_name, "RGB") for image_field in image_fields]
|
|
|
|
@torch.no_grad()
|
|
def invoke(self, context: InvocationContext) -> StringOutput:
|
|
images = self._get_images(context)
|
|
model_config = context.models.get_config(self.vllm_model)
|
|
|
|
with context.models.load(self.vllm_model).model_on_device() as (_, model):
|
|
assert isinstance(model, LlavaOnevisionForConditionalGeneration)
|
|
|
|
model_abs_path = context.models.get_absolute_path(model_config)
|
|
processor = AutoProcessor.from_pretrained(model_abs_path, local_files_only=True)
|
|
assert isinstance(processor, LlavaOnevisionProcessor)
|
|
|
|
model = LlavaOnevisionPipeline(model, processor)
|
|
output = model.run(
|
|
prompt=self.prompt,
|
|
images=images,
|
|
device=TorchDevice.choose_torch_device(),
|
|
dtype=TorchDevice.choose_torch_dtype(),
|
|
)
|
|
|
|
return StringOutput(value=output)
|