Update app.py

app.py

@@ -29,36 +29,36 @@ def summarize(text):
 
     return summary
 
-# API_URL = "https://api-inference.huggingface.co/models/deepset/bert-large-uncased-whole-word-masking-squad2"
-# headers = {"Authorization": f"Bearer {my_key}"}
-
-# def query(payload):
-# 	response = requests.post(API_URL, headers=headers, json=payload)
-# 	return response.content
-
 def answer_question(text, question):
     response = qa_model(question=question, context=text)
     answer = response['answer'] 
     return answer
 
+def query(payload):
+    API_URL = "https://api-inference.huggingface.co/models/stable-diffusion-v1-5/stable-diffusion-v1-5"
+    headers = {"Authorization": f"Bearer {my_key}"}
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+
 def summarize_and_qa(pdf_file, question):
     text = extract_text_from_pdf(pdf_file)
     summary = summarize(text)
     answer = answer_question(text, question)
-    # image_bytes = query({"inputs": answer})
-    # if image_bytes:
-    #     try:
-    #         image = Image.open(io.BytesIO(image_bytes))
-    #     except Exception as e:
-    #         return summary, answer, None
-    # else:
-    #     image = None
-    return summary, answer
+    image_bytes = query({"inputs": summary})
+    if image_bytes:
+        try:
+            image = Image.open(io.BytesIO(image_bytes))
+        except Exception as e:
+            return summary, answer, None
+    else:
+        image = None
+    
+    return summary, answer, image
 
 gr.Interface(
     fn=summarize_and_qa,
     inputs=["file", "text"],
-    outputs=["textbox", "textbox"],
+    outputs=["textbox", "textbox", "image"],
     title="Understand your PDF Better",
-    description="Upload a PDF to get a summary. You can ask any question regardging the content of the PDF. It will also generate a picture to help you better understand the content."
+    description="Upload a PDF to get a summary. You can ask any question regardging the content of the PDF. It will also generate a picture related to your PDF."
 ).launch(debug=True, share=True)