Spaces:

ybelkada
/

llava-1.5-dlai

Running on Zero

ybelkada commited on Feb 1

Commit

88f422e

•

1 Parent(s): 603d833

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -21,6 +21,6 @@ def inference(text, image):
     prompt = prompt_format.format(text)
     inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
     output = model.generate(**inputs, max_new_tokens=1024)
-    return processor.decode(output[0][2:], skip_special_tokens=True)
 gr.Interface(fn=inference, inputs=[gr.Text(), gr.Image()], outputs=gr.Text()).launch()

     prompt = prompt_format.format(text)
     inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
     output = model.generate(**inputs, max_new_tokens=1024)
+    return processor.decode(output[0], skip_special_tokens=True).split("ASSISTANT:")[-1]
 gr.Interface(fn=inference, inputs=[gr.Text(), gr.Image()], outputs=gr.Text()).launch()