Update app.py

app.py

@@ -6,10 +6,13 @@ import time
 from PIL import Image
 import torch
 import spaces
+import os
 
-processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+MODEL_ID = os.get("MODEL_ID", "llava-hf/llava-v1.6-mistral-7b-hf")
+REVISION = os.get("MODEL_REVISION", "main")
+processor = LlavaNextProcessor.from_pretrained(MODEL_ID, revision=REVISION)
 
-model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+model = LlavaNextForConditionalGeneration.from_pretrained(MODEL_ID, revision=REVISION, torch_dtype=torch.float16, low_cpu_mem_usage=True)
 model.to("cuda:0")
 
 @spaces.GPU
@@ -50,7 +53,7 @@ def bot_streaming(message, history):
     yield generated_text_without_prompt
 
 
-demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA NeXT", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]},
+demo = gr.ChatInterface(fn=bot_streaming, title="VLM Playground", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]},
                                                                       {"text": "How to make this pastry?", "files":["./baklava.png"]}], 
                         description="Try [LLaVA NeXT](https://huggingface.co/docs/transformers/main/en/model_doc/llava_next) in this demo (more specifically, the [Mistral-7B variant](https://huggingface.co/llava-hf/llava-v1.6-mistral-7b-hf)). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
                         stop_btn="Stop Generation", multimodal=True)