Update app.py

app.py

@@ -1,22 +1,26 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModel
+from PIL import Image
+import torch
 
-# Initialize the pipeline
-pipe = pipeline("feature-extraction", model="ECOFRI/CXR-LLAVA-v2", trust_remote_code=True)
+# Load the model
+model = AutoModel.from_pretrained("ECOFRI/CXR-LLAVA-v2", trust_remote_code=True)
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define a function to process inputs and return outputs
-def extract_features(text):
-    features = pipe(text)
-    return features
+# Define the function to generate the report
+def generate_report(image):
+    image = Image.open(image).convert("RGB")
+    response = model.write_radiologic_report(image)
+    return response
 
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=extract_features, 
-    inputs="text", 
-    outputs="json",
-    title="Feature Extraction Demo",
-    description="Enter text to extract features using the ECOFRI/CXR-LLAVA-v2 model."
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=generate_report,
+    inputs=gr.inputs.Image(type="file", label="Upload Chest X-ray Image"),
+    outputs=gr.outputs.Textbox(label="Radiologic Report"),
+    title="Chest X-ray Report Generator",
+    description="Upload a chest X-ray image to generate a radiologic report using the CXR-LLAVA-v2 model."
 )
 
 # Launch the interface
-iface.launch()
+interface.launch()