Update app.py

app.py

@@ -23,8 +23,6 @@ model = LlavaQwen2ForCausalLM.from_pretrained(
     attn_implementation="flash_attention_2",
     trust_remote_code=True)
 
-model.to('cuda')
-
 class KeywordsStoppingCriteria(StoppingCriteria):
     def __init__(self, keywords, tokenizer, input_ids):
         self.keywords = keywords
@@ -87,6 +85,7 @@ def bot_streaming(message, history):
         messages.append({"role": "user", "content": f"<image>\n{message['text']}"})
     elif len(history) == 0 and image is None:
         messages.append({"role": "user", "content": message['text'] })
+    model = model.to('cuda')
 
     # if image is None:
     #     gr.Error("You need to upload an image for LLaVA to work.")