Merge pull request #326 from hehaha68/add_zju_style

Update some features

Author: ly19965

app.py

@@ -164,7 +164,7 @@ def main_diffusion_inference(pos_prompt, neg_prompt,
     attr_idx = np.argmax(cnts_trigger)
     trigger_styles = ['a boy, children, ', 'a girl, children, ', 'a handsome man, ', 'a beautiful woman, ',
                       'a mature man, ', 'a mature woman, ']
-    trigger_style = '<fcsks>, ' + trigger_styles[attr_idx]
+    trigger_style = '(<fcsks>:10), ' + trigger_styles[attr_idx]
     if attr_idx == 2 or attr_idx == 4:
         neg_prompt += ', children'
 
@@ -238,7 +238,7 @@ def main_diffusion_inference_pose(pose_model_path, pose_image,
     attr_idx = np.argmax(cnts_trigger)
     trigger_styles = ['a boy, children, ', 'a girl, children, ', 'a handsome man, ', 'a beautiful woman, ',
                       'a mature man, ', 'a mature woman, ']
-    trigger_style = '<fcsks>, ' + trigger_styles[attr_idx]
+    trigger_style = '(<fcsks>:10), ' + trigger_styles[attr_idx]
     if attr_idx == 2 or attr_idx == 4:
         neg_prompt += ', children'
 
@@ -330,7 +330,7 @@ def main_diffusion_inference_multi(pose_model_path, pose_image,
     attr_idx = np.argmax(cnts_trigger)
     trigger_styles = ['a boy, children, ', 'a girl, children, ', 'a handsome man, ', 'a beautiful woman, ',
                       'a mature man, ', 'a mature woman, ']
-    trigger_style = '<fcsks>, ' + trigger_styles[attr_idx]
+    trigger_style = '(<fcsks>:10), ' + trigger_styles[attr_idx]
     if attr_idx == 2 or attr_idx == 4:
         neg_prompt += ', children'
 

facechain/constants.py

@@ -273,7 +273,7 @@ def main_diffusion_inference_inpaint(inpaint_image, strength, output_img_size, p
     attr_idx = np.argmax(cnts_trigger)
     trigger_styles = ['a boy, children, ', 'a girl, children, ', 'a handsome man, ', 'a beautiful woman, ',
                       'a mature man, ', 'a mature woman, ']
-    trigger_style = '<fcsks>, ' + trigger_styles[attr_idx]
+    trigger_style = '(<fcsks>:10), ' + trigger_styles[attr_idx]
     if attr_idx == 2 or attr_idx == 4:
         neg_prompt += ', children'
 
@@ -445,7 +445,7 @@ def main_diffusion_inference_inpaint_multi(inpaint_images, strength, output_img_
     attr_idx = np.argmax(cnts_trigger)
     trigger_styles = ['a boy, children, ', 'a girl, children, ', 'a handsome man, ', 'a beautiful woman, ',
                       'a mature man, ', 'a mature woman, ']
-    trigger_style = '<fcsks>, ' + trigger_styles[attr_idx]
+    trigger_style = '(<fcsks>:10), ' + trigger_styles[attr_idx]
     if attr_idx == 2 or attr_idx == 4:
         neg_prompt += ', children'
 

facechain/inference.py

@@ -1,23 +1,39 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
-
+import json
 from facechain.inference import GenPortrait
 import cv2
 from facechain.utils import snapshot_download
-from facechain.constants import neg_prompt, pos_prompt_with_cloth, pos_prompt_with_style, styles, base_models
+from facechain.constants import neg_prompt, pos_prompt_with_cloth, pos_prompt_with_style, base_models
 
 
 def generate_pos_prompt(style_model, prompt_cloth):
-    if style_model in base_models[0]['style_list'][:-1] or style_model is None:
-        pos_prompt = pos_prompt_with_cloth.format(prompt_cloth)
-    else:
+    if style_model is not None:
         matched = list(filter(lambda style: style_model == style['name'], styles))
         if len(matched) == 0:
             raise ValueError(f'styles not found: {style_model}')
         matched = matched[0]
-        pos_prompt = pos_prompt_with_style.format(matched['add_prompt_style'])
+        if matched['model_id'] is None:
+            pos_prompt = pos_prompt_with_cloth.format(prompt_cloth)
+        else:
+            pos_prompt = pos_prompt_with_style.format(matched['add_prompt_style'])
+    else:
+        pos_prompt = pos_prompt_with_cloth.format(prompt_cloth)
     return pos_prompt
 
+styles = []
+for base_model in base_models:
+    style_in_base = []
+    folder_path = f"/mnt/workspace/new_facechain/facechain/styles/{base_model['name']}"
+    files = os.listdir(folder_path)
+    files.sort()
+    for file in files:
+        file_path = os.path.join(folder_path, file)
+        with open(file_path, "r") as f:
+            data = json.load(f)
+            style_in_base.append(data['name'])
+            styles.append(data)
+    base_model['style_list'] = style_in_base
 
 use_main_model = True
 use_face_swap = True
@@ -28,13 +44,11 @@ def generate_pos_prompt(style_model, prompt_cloth):
 pose_image = 'poses/man/pose1.png'
 processed_dir = './processed'
 num_generate = 5
-base_model = 'ly261666/cv_portrait_model'
-revision = 'v2.0'
 multiplier_style = 0.25
 multiplier_human = 0.85
-base_model_sub_dir = 'film/film'
 train_output_dir = './output'
 output_dir = './generated'
+base_model = base_models[0]
 style = styles[0]
 model_id = style['model_id']
 
@@ -62,8 +76,8 @@ def generate_pos_prompt(style_model, prompt_cloth):
                            use_face_swap, use_post_process,
                            use_stylization)
 
-outputs = gen_portrait(processed_dir, num_generate, base_model,
-                       train_output_dir, base_model_sub_dir, revision)
+outputs = gen_portrait(processed_dir, num_generate, base_model['model_id'],
+                       train_output_dir, base_model['sub_path'], base_model['revision'])
 
 os.makedirs(output_dir, exist_ok=True)