huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 20:04:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Load later to make tests pass.

Browse Source
This commit is contained in:
Nicolas Patry 2023-08-15 14:50:35 +00:00
parent 5469316ed8
commit 4ff509948a
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
server/text_generation_server/models/idefics_causal_lm.py
View File

@ -9,7 +9,6 @@ import json
from dataclasses import dataclass
from opentelemetry import trace
from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase, ProcessorMixin
from text_generation_server.models.custom_modeling.idefics_modeling import IdeficsForVisionText2Text
from typing import Optional, Tuple, List, Type, Dict
from text_generation_server.models import Model
@ -582,6 +581,8 @@ class IdeficsCausalLM(Model):
dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False,
):
from text_generation_server.models.custom_modeling.idefics_modeling import IdeficsForVisionText2Text
if torch.cuda.is_available():
device = torch.device("cuda")
dtype = torch.float16 if dtype is None else dtype