Update app.py

app.py

@@ -6,6 +6,8 @@ from transformers import MllamaForConditionalGeneration, AutoProcessor
 from peft import PeftModel
 from huggingface_hub import login
 import spaces
+import json
+
 
 # Login to Hugging Face
 if "HF_TOKEN" not in os.environ:
@@ -20,15 +22,17 @@ processor = AutoProcessor.from_pretrained(base_model_path)
 model = MllamaForConditionalGeneration.from_pretrained(
     base_model_path,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
+    device_map="cuda",
 )
 model = PeftModel.from_pretrained(model, lora_weights_path)
+model.tie_weights()
+
 
 @spaces.GPU
-def inference(image, question):
+def inference(image):
     # Prepare input
     messages = [
-        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": question}]}
+        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe the image in JSON"}]}
     ]
     input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
     inputs = processor(image, input_text, add_special_tokens=False, return_tensors="pt").to(model.device)
@@ -39,14 +43,34 @@ def inference(image, question):
     
     # Decode output
     result = processor.decode(output[0], skip_special_tokens=True)
-    return result.strip().split("assistant\n")[1].strip()
+    json_str = result.strip().split("assistant\n")[1].strip()
+    
+    try:
+        # First JSON parse to handle escaped JSON string
+        first_parse = json.loads(json_str)
+        
+        try:
+            # Second JSON parse to get the actual JSON object
+            json_object = json.loads(first_parse)
+            # Return indented JSON string with 2 spaces
+            return json.dumps(json_object, indent=2)
+        except json.JSONDecodeError:
+            # If second parse fails, return the result of first parse indented
+            if isinstance(first_parse, (dict, list)):
+                return json.dumps(first_parse, indent=2)
+            return first_parse
+            
+    except json.JSONDecodeError:
+        # If both JSON parses fail, return original string
+        return json_str
+
+    return None  # In case of unexpected errors
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=inference,
     inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Textbox(label="Enter your question")
+        gr.Image(type="pil", label="Upload Image")
     ],
     outputs=gr.Textbox(label="Response"),
     title="Image Analysis AI",