Update app.py

app.py

@@ -27,14 +27,15 @@ model = AutoModelForCausalLM.from_pretrained("StanfordAIMI/CheXagent-8b", torch_
 
 @spaces.GPU
 def generate(image, prompt):
-    # Convert the uploaded file to an image and process
-    image = Image.open(io.BytesIO(image.read())).convert("RGB")
+    if hasattr(image_input, "read"):
+        image = Image.open(io.BytesIO(image_input.read())).convert("RGB")
+    else:
+        image = image
+
     images = [image]
 
-    # Prepare inputs
     inputs = processor(images=images, text=f" USER: <s>{prompt} ASSISTANT: <s>", return_tensors="pt").to(device=device, dtype=dtype)
 
-    # Generate the findings
     output = model.generate(**inputs, generation_config=generation_config)[0]
     response = processor.tokenizer.decode(output, skip_special_tokens=True)
     return response
@@ -51,7 +52,11 @@ with gr.Blocks() as demo:
         output_text_custom = gr.Textbox(label="Response")
 
         def custom_generate(image, prompt):
-            return generate(image, prompt)
+            if isinstance(image, str) and os.path.exists(image):
+                with open(image, 'rb') as file:
+                    return generate(file, prompt)
+            else:
+                return generate(image, prompt)
 
         generate_button_custom.click(fn=custom_generate, inputs=[image_input_custom, prompt_input_custom], outputs=output_text_custom)