Update app.py

app.py

@@ -24,19 +24,19 @@ clf_1 = pipeline(model=model_1, task="image-classification", image_processor=ima
 
 # Load the second model
 model_2_path = "Heem2/AI-vs-Real-Image-Detection"
-clf_2 = pipeline("image-classification", model=model_2_path)
+clf_2 = pipeline("image-classification", model=model_2_path, device=device)
 
 # Load additional models
 models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
 
 # Load the third and fourth models
-feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0])
+feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0], device=device)
 model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
 
-feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1])
+feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1], device=device)
 model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
 
-model_5 = load_model("large_model_3lakh_v1.h5")
+model_5 = load_model("large_model_3lakh_v1.h5", device=device)
 # Define class names for all models
 class_names_1 = ['artificial', 'real']
 class_names_2 = ['AI Image', 'Real Image']