add text reponse

app.py

@@ -50,14 +50,14 @@ with gr.Blocks(title='PQ3D Demo') as demo:
     
     def inference_wrapper(text):
         scan_id = model_3d.value['orig_name'].split('.')[0]
-        inst_id = inference(scan_id, text)
-        return f"assets/mask/{scan_id}/{scan_id}_obj_{inst_id}.glb"
+        inst_id, response = inference(scan_id, text)
+        return f"assets/mask/{scan_id}/{scan_id}_obj_{inst_id}.glb", response
     
     gr.Interface(
         fn=inference_wrapper,
         inputs=["text"],
-        outputs=gr.Model3D(
-                clear_color=[0.0, 0.0, 0.0, 0.0],  camera_position=(80, 100, 6), label="3D Model"),
+        outputs=[gr.Model3D(
+                clear_color=[0.0, 0.0, 0.0, 0.0],  camera_position=(80, 100, 6), label="3D Model"), "text"],
         examples=[
             ["armchair"], ["Sofa"], ["left computer on the desk"]
         ],

pq3d/inference.py

@@ -171,9 +171,14 @@ def inference(scan_id, text):
     model = Query3DUnified()
     load_msg = model.load_state_dict(torch.load(os.path.join(CKPT_DIR, 'pytorch_model.bin'), map_location='cpu'), strict=False)
     data_dict = model(data_dict)
+    # calculate result id
     result_id = data_dict['obj_ids'][0][torch.argmax(data_dict['og3d_logits'][0]).item()]
     print(f"finish infernece result id is {result_id}")
-    return result_id
+    # calculate langauge
+    tokenizer = AutoTokenizer.from_pretrained("t5-small")
+    response_pred = tokenizer.batch_decode(data_dict['generation_logits'], skip_special_tokens=True)[0]
+    print(f"text response is {response_pred}")
+    return result_id, response_pred
     
 if __name__ == '__main__':
     inference("scene0050_00", "chair")