.

scripts/generate_prompt.py

@@ -7,6 +7,20 @@ from huggingface_hub import hf_hub_download
 
 import spaces
 
+import tensorflow as tf
+
+gpus = tf.config.list_physical_devices('GPU')
+if gpus:
+    try:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        print("GPUs are available and memory growth is set.")
+    except RuntimeError as e:
+        print("Error setting GPU memory growth:", e)
+else:
+    print("No GPUs are available.")
+
+
 # 画像サイズの設定
 IMAGE_SIZE = 448