Update app.py

app.py

@@ -12,10 +12,6 @@ from PIL import UnidentifiedImageError, Image
 import gradio as gr
 import matplotlib.pyplot as plt
 
-# Ensure TensorFlow uses GPU
-print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
-assert len(tf.config.list_physical_devices('GPU')) > 0, "No GPU available!"
-
 # Load the dataset
 dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
 
@@ -27,8 +23,22 @@ dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
 image_dir = 'civitai_images'
 os.makedirs(image_dir, exist_ok=True)
 
+# Try to use GPU, fall back to CPU if not available
+try:
+    gpus = tf.config.list_physical_devices('GPU')
+    if gpus:
+        tf.config.experimental.set_memory_growth(gpus[0], True)
+        device = '/GPU:0'
+        print("Using GPU")
+    else:
+        raise RuntimeError("No GPU found")
+except RuntimeError as e:
+    print(e)
+    device = '/CPU:0'
+    print("Using CPU")
+
 # Load the ResNet50 model pretrained on ImageNet
-with tf.device('/GPU:0'):
+with tf.device(device):
     model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
 
 # Function to extract features
@@ -37,7 +47,7 @@ def extract_features(img_path, model):
     img_array = image.img_to_array(img)
     img_array = np.expand_dims(img_array, axis=0)
     img_array = preprocess_input(img_array)
-    with tf.device('/GPU:0'):
+    with tf.device(device):
         features = model.predict(img_array)
     return features.flatten()
 
@@ -136,4 +146,4 @@ interface = gr.Interface(
     description="Upload an image and get 5 recommended similar images with model names and distances."
 )
 
-interface.launch()
+interface.launch()