download the model and processor on startup

app.py

@@ -1,17 +1,22 @@
 import gradio as gr
 from predict import prediction
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from functools import partial
 
 title = "Audio deepfake Classification"
 description = """This space uses fine-tuned kubinooo/convnext-tiny-224-audio-deepfake-classification model to classify audio files.
 """
 
+processor = AutoImageProcessor.from_pretrained("kubinooo/convnext-tiny-224-audio-deepfake-classification")
+model = AutoModelForImageClassification.from_pretrained("kubinooo/convnext-tiny-224-audio-deepfake-classification")
+
 demo = gr.Interface(
     title=title,
     inputs=gr.Audio( type="filepath",
         interactive=True,  # This prevents users from uploading their own images
         show_label=True,
         label="Select from examples below or upload/record your own audio"),
-    fn=prediction,
+    fn=partial(prediction, processor=processor, model=model),
     outputs=gr.Label(
         num_top_classes=2,
     ),
@@ -21,5 +26,4 @@ demo = gr.Interface(
     ],
     description=description
 )
-
 demo.launch(share=True,debug=True)

predict.py

@@ -1,11 +1,7 @@
 import torch
-from transformers import AutoImageProcessor, AutoModelForImageClassification
 from process_audio import create_mel_spectrograms
-import os
 
-def predict_image(image):
-  processor = AutoImageProcessor.from_pretrained("kubinooo/convnext-tiny-224-audio-deepfake-classification")
-  model = AutoModelForImageClassification.from_pretrained("kubinooo/convnext-tiny-224-audio-deepfake-classification")
+def predict_image(image, processor, model):
 
   if image.mode != 'RGB':
     image = image.convert('RGB')
@@ -30,14 +26,14 @@ def predict_image(image):
     return {"real": 0.0, "fake": 1.0}
 
 
-def prediction(file_path):
+def prediction(file_path, processor, model):
   total_real = 0.0
   total_fake = 0.0
 
   pil_images = create_mel_spectrograms(file_path, 2, 0)
 
   for image in pil_images:
-    pred = predict_image(image)
+    pred = predict_image(image, processor, model)
     total_real += pred["real"]
     total_fake += pred["fake"]