Attempt to add image model

app.py

@@ -1,14 +1,21 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the model using pipeline
-pipe = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
+# Load the models using pipeline
+audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
+image_model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")
 
 # Define the prediction function
-def predict(audio):
-    print("Audio file received:", audio)  # Debugging statement
+def predict(data, model_choice):
+    print("Data received:", data)  # Debugging statement
     try:
-        result = pipe(audio)
+        if model_choice == "Audio Deepfake Detection":
+            result = audio_model(data)
+        elif model_choice == "Image Deepfake Detection":
+            result = image_model(data)
+        else:
+            return {"error": "Invalid model choice"}
+        
         print("Raw prediction result:", result)  # Debugging statement
         # Convert the result to the expected format
         output = {item['label']: item['score'] for item in result}
@@ -18,14 +25,34 @@ def predict(audio):
         print("Error during prediction:", e)  # Debugging statement
         return {"error": str(e)}
 
+# Define the interface based on the selected model
+def update_interface(model_choice):
+    if model_choice == "Audio Deepfake Detection":
+        return gr.Audio(type="filepath"), gr.Label()
+    elif model_choice == "Image Deepfake Detection":
+        return gr.Image(type="filepath"), gr.Label()
+    else:
+        return None, None
+
 # Create the Gradio interface
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Audio(type="filepath"),
-    outputs=gr.Label(),
-    title="Testing Deepfake Audio Detection Simple Interface",
-    description="Upload an audio file or record your voice to detect if the audio is a deepfake."
-)
+with gr.Blocks() as iface:
+    model_choice = gr.Radio(choices=["Audio Deepfake Detection", "Image Deepfake Detection"], label="Select Model", value="Audio Deepfake Detection")
+    input_component, output_component = update_interface(model_choice.value)
+    
+    def update_inputs(model_choice):
+        input_component, output_component = update_interface(model_choice)
+        input_placeholder.update(visible=False)
+        output_placeholder.update(visible=False)
+        input_placeholder.update(visible=True, component=input_component)
+        output_placeholder.update(visible=True, component=output_component)
+        
+    input_placeholder = gr.Placeholder(gr.Component, visible=True)
+    output_placeholder = gr.Placeholder(gr.Component, visible=True)
+    
+    model_choice.change(fn=update_inputs, inputs=model_choice, outputs=[input_placeholder, output_placeholder])
+    
+    submit_button = gr.Button("Submit")
+    submit_button.click(fn=predict, inputs=[input_placeholder, model_choice], outputs=output_placeholder)
+    
+    iface.launch()
 
-# Launch the interface
-iface.launch()