Update app.py

app.py

@@ -1,3 +1,55 @@
 import gradio as gr
+from PIL import Image
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-gr.load("models/ManishThota/Thota").launch()
+#Setting device to cuda
+torch.set_default_device("cuda")
+
+# # Ensure GPU usage if available
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# torch.set_default_tensor_type('torch.cuda.FloatTensor' if device=='cuda' else 'torch.FloatTensor')
+
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+
+
+# Initialize the model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("ManishThota/Thota",
+                                             torch_dtype=torch.float16, 
+                                             device_map="auto",
+                                             trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("ManishThota/Thota", trust_remote_code=True)
+
+def predict_answer(question, max_tokens):
+    #Set inputs
+    text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{question}? ASSISTANT:"
+    image = image.convert("RGB")
+    
+    input_ids = tokenizer(text, return_tensors='pt').input_ids.to('cuda')
+    
+    #Generate the answer
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=max_tokens,
+        images=image_tensor,
+        use_cache=True)[0]
+    
+    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
+
+def gradio_predict(question, max_tokens):
+    answer = predict_answer(question, max_tokens)
+    return answer
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=gradio_predict,
+    inputs=[ gr.Textbox(label="Question", placeholder="e.g. Who is Manish Kumar Thota?", scale=4),
+            gr.Slider(2, 100, value=25, label="Token Count", info="Choose between 2 and 100")],
+    outputs=gr.TextArea(label="Answer"),
+    title="Welcome to Manish Kumar Thota's Portfolio",
+    description="Ask questions about Manish and get to know him as if you are directly talking to him, this model is still in development so your feedback is highly valuable to improve the model.",
+)
+
+# Launch the app
+iface.queue().launch(debug=True)