|
import gradio as gr |
|
import spaces |
|
from PIL import Image |
|
from models.mllava import MLlavaProcessor, LlavaForConditionalGeneration, chat_mllava, MLlavaForConditionalGeneration |
|
from typing import List |
|
processor = MLlavaProcessor() |
|
model = LlavaForConditionalGeneration.from_pretrained("MFuyu/mllava_v2_4096") |
|
|
|
@spaces.GPU |
|
def generate(text:str, images:List[Image.Image], history: List[dict]): |
|
model = model.to("cuda") |
|
|
|
for text, history in chat_mllava(text, images, model, processor, history=history, stream=True): |
|
yield text, history |
|
|
|
def build_demo(): |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
processor = MLlavaProcessor() |
|
model = LlavaForConditionalGeneration.from_pretrained("MFuyu/mllava_v2_4096") |
|
demo = build_demo() |
|
demo.launch() |
