arg flag追加

app.py

@@ -3,6 +3,7 @@ from flask_socketio import SocketIO, emit
 from flask_cors import CORS
 import io
 import os
+import argparse
 from PIL import Image
 import torch
 import gc
@@ -183,5 +184,10 @@ def server_error(e):
     return jsonify(error=str(e)), 500
 
 if __name__ == '__main__':
-    initialize(local_model=True)
+    parser = argparse.ArgumentParser(description='Server options.')
+    parser.add_argument('--local_model', type=bool, default=False, help='Use local model')
+    parser.add_argument('--use_gpu', type=bool, default=True, help='Set to True to use GPU but if not available, it will use CPU')
+    args = parser.parse_args()
+    
+    initialize(local_model=args.local_model, use_gpu=args.use_gpu)
     socketio.run(app, debug=True, host='0.0.0.0', port=5000)

process_utils.py

@@ -18,9 +18,8 @@ load_dotenv()
 # グローバル変数
 local_model = False
 model = None
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-device = "cpu"
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
+device = None
+torch_dtype = None # torch.float16 if device == "cuda" else torch.float32
 sotai_gen_pipe = None
 refine_gen_pipe = None
 
@@ -44,9 +43,10 @@ def ensure_rgb(image):
         return image.convert('RGB')
     return image
 
-def initialize(_local_model=False):
-    global model, sotai_gen_pipe, refine_gen_pipe, local_model
-
+def initialize(_local_model=False, use_gpu=True)
+    global model, sotai_gen_pipe, refine_gen_pipe, local_model, device, torch_dtype
+    device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
     local_model = _local_model
     model = load_wd14_tagger_model()
     sotai_gen_pipe = initialize_sotai_model()
@@ -225,7 +225,7 @@ def generate_sotai_image(input_image: Image.Image, output_width: int, output_hei
             image=[input_image, input_image],
             negative_prompt=f"(wings:1.6), (clothes, garment, lighting, gray, missing limb, extra line, extra limb, extra arm, extra legs, hair, bangs, fringe, forelock, front hair, fill:1.4), (ink pool:1.6)",
             # negative_prompt=f"{easy_negative_v2}, (wings:1.6), (clothes, garment, lighting, gray, missing limb, extra line, extra limb, extra arm, extra legs, hair, bangs, fringe, forelock, front hair, fill:1.4), (ink pool:1.6)",
-            num_inference_steps=40,
+            num_inference_steps=20,
             guidance_scale=8,
             width=output_width,
             height=output_height,