Update app.py

app.py

@@ -1,79 +1,71 @@
-import gradio as gr
-import google.generativeai as genai
 import os
-import io
+import time
 import mimetypes
-from PIL import Image
+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
 
-# Initialize Gemini API
-try:
-    GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
-    genai.configure(api_key=GOOGLE_API_KEY)
-    model = genai.GenerativeModel("gemini-pro-vision")
-except KeyError:
-    print("Error: GEMINI_API_KEY not found in environment variables. Please set it in Hugging Face Space secrets.")
-    exit()
-except Exception as e:
-    print(f"Error initializing Gemini API: {e}")
-    exit()
+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 generate_response(text_input, file_input):
-    """Generates a response using Gemini's multiprompt feature with various file types."""
-    try:
-        contents = []
-        if text_input:
-            contents.append(text_input)
+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]))
 
-        if file_input:
-            file_path = file_input
-            mime_type, _ = mimetypes.guess_type(file_path)
-            if mime_type is None:
-                return "Unsupported file type. Please upload a valid image, video, audio, or PDF."
-            
-            try:
-                with open(file_path, "rb") as f:
-                    file_bytes = f.read()
-            except Exception as e:
-                return f"Error reading file: {e}"
-            
-            if mime_type.startswith("image"):
-                try:
-                    # Check if it's a valid image
-                    Image.open(io.BytesIO(file_bytes))
-                    contents.append({
-                        "mime_type": mime_type,
-                        "data": file_bytes
-                    })
-                except Exception as e:
-                     return f"Error processing image: {e}"
-            elif mime_type.startswith("video") or mime_type.startswith("audio") or mime_type == "application/pdf":
-                contents.append({
-                    "mime_type": mime_type,
-                    "data": file_bytes
-                })
-            else:
-                return "Unsupported file type. Please upload a valid image, video, audio, or PDF."
+    # 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
 
-        if not contents:
-            return "Please provide either text or a file."
+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")
 
-        response = model.generate_content(contents)
-        response.resolve()
-        return response.text
-    except Exception as e:
-        return f"Error generating response: {e}"
+        submit_button.click(fn=handle_files, inputs=[file_input], outputs=[output])
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=generate_response,
-    inputs=[
-        gr.Textbox(lines=3, placeholder="Enter text prompt here..."),
-        gr.File(label="Upload a File (Image, Video, Audio, PDF)")
-    ],
-    outputs=gr.Textbox(label="Gemini Response"),
-    title="Multimodal Gemini Chatbot",
-    description="Upload an image, video, audio, or PDF, or enter text to interact with Gemini. You can provide both text and a file for multiprompt."
-)
+    return app
 
-if __name__ == "__main__":
-    iface.launch(share=False)
+app = gradio_interface()
+app.launch()