use_local編集

app.py

@@ -79,7 +79,7 @@ def worker():
             if task.task_id in active_tasks:
                 future = executor.submit(process_task, task)
                 task_futures[task.task_id] = future
-            update_queue_status(f'Task started: {task.task_id}')
+            update_queue_status(f'Task processing: {task.task_id}')
         except Exception as e:
             print(f"Worker error: {str(e)}")
         finally:
@@ -187,7 +187,8 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Server options.')
     parser.add_argument('--use_local', action='store_true', help='Use local model')
     parser.add_argument('--use_gpu', action='store_true', help='Set to True to use GPU but if not available, it will use CPU')
+    parser.add_argument('--use_dotenv', action='store_true', help='Use .env file for environment variables')
     args = parser.parse_args()
     
-    initialize(args.use_local, args.use_gpu)
-    socketio.run(app, debug=True, host='0.0.0.0', port=80)
+    initialize(args.use_local, args.use_gpu, args.use_dotenv)
+    socketio.run(app, debug=True, host='0.0.0.0', port=5000)

scripts/anime.py

@@ -6,15 +6,22 @@ Example:
 
 import os
 import torch
-from scripts.data import get_image_list, get_transform
+from scripts.data import get_image_list, get_transform, tensor_to_img, save_image
 from scripts.model import create_model
-from scripts.data import tensor_to_img, save_image
 import argparse
 from tqdm.auto import tqdm
 from kornia.enhance import equalize_clahe
 from PIL import Image
 import numpy as np
 
+model = None
+
+def init_model(use_local=False):
+    global model
+    model_opt = "default"
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = create_model(model_opt, use_local).to(device)
+    model.eval()
 
 # numpy配列の画像を受け取り、線画を生成してnumpy配列で返す
 def generate_sketch(image, clahe_clip=-1, load_size=512):
@@ -28,10 +35,10 @@ def generate_sketch(image, clahe_clip=-1, load_size=512):
         np.ndarray: output image
     """
     # create model
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model_opt = "default"
-    model = create_model(model_opt).to(device)
-    model.eval()
+    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # model_opt = "default"
+    # model = create_model(model_opt).to(device)
+    # model.eval()
     
     aus_resize = None
     if load_size > 0:
@@ -85,7 +92,7 @@ if __name__ == '__main__':
     # create model
     gpu_list = ','.join(str(x) for x in opt.gpu_ids)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = create_model(opt.model).to(device)      # create a model given opt.model and other options
+    model = create_model(opt.model, use_local=True).to(device)      # create a model given opt.model and other options
     model.eval()
     
     for test_path in tqdm(get_image_list(opt.dataroot)):

scripts/model.py

@@ -144,7 +144,7 @@ class Upsample(nn.Module):
         return x
 
 
-def create_model(model):
+def create_model(model, use_local):
     """Create a model for anime2sketch
     hardcoding the options for simplicity
     """
@@ -156,7 +156,10 @@ def create_model(model):
     cwd = os.getcwd() # 現在のディレクトリを保存
     os.chdir(os.path.dirname(__file__)) # このファイルのディレクトリに移動
     if model == 'default':
-        model_path = download_file("netG.pth", subfolder="models/Anime2Sketch")
+        model_path = (lambda filename, subfolder: os.path.join(subfolder, filename) if use_local else download_file(filename, subfolder)) \
+                        ("netG.pth", "models/Anime2Sketch")
+        # model_path = ((filename, subfolder) => if (use_local) os.path.join(subfolder, filename) else download_file(filename, subfolder))("netG.pth", "models/Anime2Sketch") // JavaScript
+
         ckpt = torch.load(model_path)
         for key in list(ckpt.keys()):
             if 'module.' in key:

scripts/process_utils.py

@@ -6,6 +6,7 @@ import cv2
 import numpy as np
 from scripts.generate_prompt import load_wd14_tagger_model, generate_tags, preprocess_image as wd14_preprocess_image
 from scripts.lineart_util import scribble_xdog, get_sketch, canny
+from scripts.anime import init_model
 import torch
 from diffusers import StableDiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler, AutoencoderKL
 import gc
@@ -32,8 +33,9 @@ def ensure_rgb(image):
         return image.convert('RGB')
     return image
 
-def initialize(_use_local, use_gpu):
-    # load_dotenv()
+def initialize(_use_local=False, use_gpu=False, use_dotenv=False):
+    if use_dotenv:
+        load_dotenv()
     global model, sotai_gen_pipe, refine_gen_pipe, use_local, 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
@@ -41,6 +43,7 @@ def initialize(_use_local, use_gpu):
     print('')
     print(f"Device: {device}, Local model: {_use_local}")
     print('')
+    init_model(use_local)
     model = load_wd14_tagger_model()
     sotai_gen_pipe = initialize_sotai_model()
     refine_gen_pipe = initialize_refine_model()
@@ -53,8 +56,8 @@ def initialize_sotai_model():
     global device, torch_dtype
 
     sotai_sd_model_path = get_file_path(os.environ["sotai_sd_model_name"], subfolder=os.environ["sd_models_dir"])
-    # controlnet_path1 =  get_file_path(os.environ["controlnet_name1"], subfolder=os.environ["controlnet_dir2"])
-    controlnet_path1 =  get_file_path(os.environ["controlnet_name2"], subfolder=os.environ["controlnet_dir1"])
+    controlnet_path1 =  get_file_path(os.environ["controlnet_name1"], subfolder=os.environ["controlnet_dir2"])
+    # controlnet_path1 =  get_file_path(os.environ["controlnet_name2"], subfolder=os.environ["controlnet_dir1"])
     controlnet_path2 =  get_file_path(os.environ["controlnet_name2"], subfolder=os.environ["controlnet_dir1"])
     print(use_local, controlnet_path1)
 
@@ -229,7 +232,7 @@ def generate_sotai_image(input_image: Image.Image, output_width: int, output_hei
             denoising_strength=0.13,
             num_images_per_prompt=1,  # Equivalent to batch_size
             guess_mode=[True, True],  # Equivalent to pixel_perfect
-            controlnet_conditioning_scale=[1.2, 1.3],  # 各ControlNetの重み
+            controlnet_conditioning_scale=[1.4, 1.3],  # 各ControlNetの重み
             guidance_start=[0.0, 0.0],
             guidance_end=[1.0, 1.0],
         )