Llava Class config updated

app.py

@@ -1,48 +1,43 @@
-# import gradio as gr
-#
-# def greet(name):
-#     return "Hello " + name + "!!"
-#
-# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-# demo.launch()
-
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
 import torch
 
-# Load LLaVA model
 model_id = "llava-hf/llava-1.5-7b-hf"
 processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id,
-                                             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                                             device_map="auto")
-
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    low_cpu_mem_usage=True,
+    device_map="auto",
+)
 
 def chat_with_llava(image, question, history=[]):
     if image is None or not question.strip():
-        return history + [[question, "Please provide both an image and a prompt."]]
-
-    # Preprocess input
-    inputs = processor(text=question, images=image, return_tensors="pt").to(model.device)
+        history.append([question, "Please provide both an image and a question."])
+        return history
+
+    # Format multimodal prompt
+    conversation = [
+        {"role": "user", "content": [
+            {"type": "text", "text": question},
+            {"type": "image"}
+        ]}
+    ]
+    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+    # Encode inputs
+    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=512)
+    answer = processor.decode(outputs[0], skip_special_tokens=True)
 
-    # Generate output
-    output = model.generate(**inputs, max_new_tokens=512)
-    answer = processor.decode(output[0], skip_special_tokens=True)
-
-    # Append to chat history
     history.append([question, answer])
     return history
 
-
-# Create Gradio chatbot interface
 chat_interface = gr.ChatInterface(
     fn=chat_with_llava,
-    inputs=[gr.Image(type="pil", label="Palm Image"),
-            gr.Textbox(label="Your Question", placeholder="What does my palm say about my future?")],
+    inputs=[gr.Image(type="pil", label="Palm Image"), gr.Textbox(label="Your Question")],
     title="🖐️ AI Palm Reader",
-    description="Upload a palm image and ask a question. The LLaVA model will respond like a palmistry expert!",
+    description="Upload your palm image and ask a question—LLaVA will respond with a palmistry-style reading."
 )
-
 chat_interface.launch()
-