Rahatara's picture
Update app.py
f406088 verified
raw
history blame
2.53 kB
import os
import time
import mimetypes
import google.generativeai as genai
import gradio as gr
# Configure the Gemini API
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
def upload_to_gemini(file_path, mime_type=None):
"""Uploads the given file to Gemini, returning the URI for use in chat."""
file = genai.upload_file(file_path, mime_type=mime_type)
print(f"Uploaded file '{file.display_name}' as: {file.uri}")
return file.uri
def wait_for_files_active(file_uris):
"""Ensure all files are processed and ready for use."""
print("Waiting for file processing...")
for uri in file_uris:
while True:
file = genai.get_file(uri)
if file.state.name == "ACTIVE":
print("File ready.")
break
elif file.state.name == "FAILED":
raise Exception(f"File {uri} failed to process")
time.sleep(10)
def handle_files(files):
"""Handle file uploads, process them, and interact with Gemini API."""
# Save uploaded files locally
saved_files = []
for file in files:
local_path = f"/tmp/{file.name}"
with open(local_path, "wb") as f:
f.write(file.read())
saved_files.append((local_path, mimetypes.guess_type(local_path)[0]))
# Upload files to Gemini and wait for them to be ready
file_uris = [upload_to_gemini(path, mime_type) for path, mime_type in saved_files]
wait_for_files_active(file_uris)
# Start a chat session with uploaded files
model = genai.GenerativeModel(
model_name="gemini-2.0-flash-exp",
generation_config={
"temperature": 0.4,
"top_p": 0.95,
"top_k": 40,
"max_output_tokens": 8192,
"response_mime_type": "text/plain",
}
)
chat_session = model.start_chat(history=[{"role": "user", "parts": [{"file_uri": uri} for uri in file_uris]}])
response = chat_session.send_message("Describe the contents.")
return response.text
def gradio_interface():
"""Setup Gradio UI for file uploads and processing."""
with gr.Blocks() as app:
with gr.Row():
file_input = gr.File(label="Upload files", file_types=[".pdf", ".png", ".jpg", ".mp3", ".mp4"], multiple=True)
submit_button = gr.Button("Submit")
output = gr.Textbox(label="Response from Gemini")
submit_button.click(fn=handle_files, inputs=[file_input], outputs=[output])
return app
app = gradio_interface()
app.launch()