.

scripts/generate_prompt.py

@@ -2,15 +2,11 @@ import csv
 import os
 import cv2
 import numpy as np
-import tensorflow as tf
 from tensorflow.keras.layers import TFSMLayer
 from huggingface_hub import hf_hub_download
 
 import spaces
 
-# TensorFlowがGPUを使用しないように設定
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
 # 画像サイズの設定
 IMAGE_SIZE = 448
 
@@ -40,7 +36,6 @@ def download_model_files(repo_id, model_dir, sub_dir, files, sub_files):
     for file in sub_files:
         hf_hub_download(repo_id, file, subfolder=sub_dir, cache_dir=os.path.join(model_dir, sub_dir), force_download=True, force_filename=file)
 
-@spaces.GPU
 def load_wd14_tagger_model():
     """WD14タグ付けモデルをロード"""
     model_dir = "wd14_tagger_model"

scripts/process_utils.py

@@ -43,7 +43,7 @@ def initialize(_use_local=False, use_gpu=False, use_dotenv=False):
     print(f"\nDevice: {device}, Local model: {_use_local}\n")
 
     init_model(use_local)
-    model = load_wd14_tagger_model()
+    # model = load_wd14_tagger_model()
     sotai_gen_pipe = initialize_sotai_model()
     refine_gen_pipe = initialize_refine_model()
 
@@ -154,8 +154,7 @@ def initialize_refine_model():
 def get_wd_tags(images: list) -> list:
     global model
     if model is None:
-        raise ValueError("Model is not initialized")
-        # initialize()
+        model = load_wd14_tagger_model()
     preprocessed_images = [wd14_preprocess_image(img) for img in images]
     preprocessed_images = np.array(preprocessed_images)
     return generate_tags(preprocessed_images, os.environ["wd_model_name"], model)