diff --git a/src/backend/app/database/__init__.py b/src/backend/app/database/__init__.py index e7fe003..990816e 100644 --- a/src/backend/app/database/__init__.py +++ b/src/backend/app/database/__init__.py @@ -112,7 +112,9 @@ class DataType(db.Model): data_type_code = db.Column(String(50), nullable=False) instance_prompt = db.Column(Text, comment='数据集相关的Prompt (Instance Prompt Template, e.g. "a photo of sks person")') class_prompt = db.Column(String(255), comment='类别Prompt (e.g. "a photo of person")') - placeholder_token = db.Column(String(50), comment='TI Placeholder (e.g. "")') + validation_prompt_prefix_db_lora = db.Column(Text, comment='DreamBooth/LoRA验证生成图Prompt前缀') + validation_prompt_prefix_ti = db.Column(Text, comment='Textual Inversion验证生成图Prompt前缀') + placeholder_token = db.Column(String(50), comment='TI Placeholder (e.g. "")') initializer_token = db.Column(String(50), comment='TI Initializer (e.g. "person")') description = db.Column(Text) diff --git a/src/backend/app/workers/finetune_worker.py b/src/backend/app/workers/finetune_worker.py index 8d66f8e..d91b7c6 100644 --- a/src/backend/app/workers/finetune_worker.py +++ b/src/backend/app/workers/finetune_worker.py @@ -71,7 +71,9 @@ def run_finetune_task(task_id, finetune_method, train_images_dir, # 默认值 (Fallback) instance_prompt = "a photo of sks person" class_prompt = "a photo of person" - placeholder_token = "" + validation_prompt_prefix_db_lora = "((a selfie photo of person face)), distinct feature, head and shoulders shot, front view, face looking at camera" + validation_prompt_prefix_ti = "a selfie photo of person" + placeholder_token = "" initializer_token = "person" if data_type: @@ -79,30 +81,30 @@ def run_finetune_task(task_id, finetune_method, train_images_dir, instance_prompt = data_type.instance_prompt if data_type.class_prompt: class_prompt = data_type.class_prompt + if data_type.validation_prompt_prefix_db_lora: + validation_prompt_prefix_db_lora = data_type.validation_prompt_prefix_db_lora + if data_type.validation_prompt_prefix_ti: + validation_prompt_prefix_ti = data_type.validation_prompt_prefix_ti if data_type.placeholder_token: placeholder_token = data_type.placeholder_token if data_type.initializer_token: initializer_token = data_type.initializer_token - logger.info(f"DataType Config - Template: '{instance_prompt}', Class: '{class_prompt}'") - - # 根据微调方法调整 Instance Prompt + # 根据微调方法选择对应的 validation_prompt_prefix if finetune_method == 'textual_inversion': - # TI: 将 'sks' 替换为 placeholder_token - instance_prompt_prefix = instance_prompt.replace('sks', placeholder_token) - else: - # DreamBooth/LoRA: 直接使用模板 - instance_prompt_prefix = instance_prompt + validation_prompt_prefix = validation_prompt_prefix_ti + else: # dreambooth 或 lora + validation_prompt_prefix = validation_prompt_prefix_db_lora + + logger.info(f"DataType Config - Instance: '{instance_prompt}', Class: '{class_prompt}', Validation Prefix ({finetune_method}): '{validation_prompt_prefix}'") - # 处理 Validation Prompt + # 处理 Validation Prompt：使用 validation_prompt_prefix + custom_prompt prompt_suffix = finetune.custom_prompt.strip() if finetune.custom_prompt else "" if prompt_suffix: - validation_prompt = f"{instance_prompt_prefix}, {prompt_suffix}" + validation_prompt = f"{validation_prompt_prefix}, {prompt_suffix}" else: - validation_prompt = instance_prompt_prefix - - instance_prompt = instance_prompt_prefix + validation_prompt = validation_prompt_prefix logger.info(f"Prompts Finalized - Instance: '{instance_prompt}', Class: '{class_prompt}', Validation: '{validation_prompt}'") diff --git a/src/backend/app/workers/heatmap_worker.py b/src/backend/app/workers/heatmap_worker.py index f5081e0..0d57dd5 100644 --- a/src/backend/app/workers/heatmap_worker.py +++ b/src/backend/app/workers/heatmap_worker.py @@ -60,7 +60,7 @@ def run_heatmap_task(task_id, original_image_path, perturbed_image_path, logger.info(f"Starting heatmap task {task_id}") # 从数据库获取提示词（从关联的Perturbation任务获取） - prompt_text = "a photo of sks person" # 默认值 + prompt_text = "a selfie photo of person" # 默认值 target_word = "person" # 默认值 # 通过flow_id查找关联的Perturbation任务 diff --git a/src/backend/app/workers/perturbation_worker.py b/src/backend/app/workers/perturbation_worker.py index 792d685..10ae359 100644 --- a/src/backend/app/workers/perturbation_worker.py +++ b/src/backend/app/workers/perturbation_worker.py @@ -109,6 +109,11 @@ def run_perturbation_task(task_id, algorithm_code, epsilon, input_dir, output_di # 保存扰动图片到数据库 _save_perturbed_images(task_id, output_dir) + + logs_dir = os.path.join(output_dir, 'logs') + if os.path.exists(logs_dir) and os.path.isdir(logs_dir): + logger.info(f"Final cleanup of logs directory: {logs_dir}") + shutil.rmtree(logs_dir) # 更新任务状态为完成 completed_status = TaskStatus.query.filter_by(task_status_code='completed').first() diff --git a/src/backend/config/algorithm_config.py b/src/backend/config/algorithm_config.py index 6d89d3c..53662d5 100644 --- a/src/backend/config/algorithm_config.py +++ b/src/backend/config/algorithm_config.py @@ -46,12 +46,12 @@ class AlgorithmConfig: }, 'picasso': { 'name': '毕加索立体派', - 'prompt': 'cubist painting by picasso', + 'prompt': 'cubism painting by picasso', 'description': '模仿毕加索的立体主义风格' }, 'baroque': { 'name': '巴洛克风格', - 'prompt': 'baroque style painting', + 'prompt': 'oil painting in baroque style', 'description': '经典巴洛克艺术风格' } } @@ -91,8 +91,8 @@ class AlgorithmConfig: 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model1'], 'enable_xformers_memory_efficient_attention': True, - 'instance_prompt': 'a photo of sks person', - 'class_prompt': 'a photo of person', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of person', 'num_class_images': 200, 'center_crop': True, 'with_prior_preservation': True, @@ -115,8 +115,8 @@ class AlgorithmConfig: 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model1'], 'enable_xformers_memory_efficient_attention': True, - 'instance_prompt': 'a photo of person', - 'class_prompt': 'a photo of person', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of person', 'num_class_images': 100, 'center_crop': True, 'with_prior_preservation': True, @@ -138,7 +138,7 @@ class AlgorithmConfig: 'conda_env': CONDA_ENVS['caat'], 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model2'], - 'instance_prompt': 'a photo of a person', + 'instance_prompt': 'a painting in style', 'resolution': 512, 'learning_rate': 1e-5, 'lr_warmup_steps': 0, @@ -155,8 +155,8 @@ class AlgorithmConfig: 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model2'], 'with_prior_preservation': True, - 'instance_prompt': 'a photo of a person', - 'class_prompt': 'person', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of person', 'num_class_images': 200, 'resolution': 512, 'learning_rate': 1e-5, @@ -191,7 +191,7 @@ class AlgorithmConfig: 'center_crop': True, 'max_train_steps': 150, 'eps': 0.05, - 'target_style': 'cubism painting by picasso', + 'target_style': 'impressionism painting by van gogh', 'style_strength': 0.75, 'n_runs': 3, 'style_transfer_iter': 15, @@ -206,8 +206,8 @@ class AlgorithmConfig: 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model1'], 'enable_xformers_memory_efficient_attention': True, - 'instance_prompt': 'a photo of sks person', - 'class_prompt': 'a photo of person', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of persons', 'num_class_images': 100, 'center_crop': True, 'with_prior_preservation': True, @@ -324,7 +324,7 @@ class AlgorithmConfig: 'resolution': 512, 'train_batch_size': 1, 'gradient_accumulation_steps': 1, - 'learning_rate': 5e-7, + 'learning_rate': 1e-6, 'lr_scheduler': 'constant_with_warmup', 'lr_warmup_steps': 50, 'num_class_images': 100, @@ -334,7 +334,7 @@ class AlgorithmConfig: 'mixed_precision': 'bf16', 'prior_generation_precision': 'bf16', 'sample_batch_size': 5, - 'validation_prompt': 'a selfie photo of person, head-and-shoulders, face looking at the camera, Eiffel Tower clearly visible behind, outdoor daytime, realistic', + 'validation_prompt': '((a selfie photo of person face)), distinct feature, head and shoulders shot, front view, face looking at camera', 'num_validation_images': 5, 'validation_num_inference_steps': 120, 'validation_guidance_scale': 7.0, @@ -349,8 +349,8 @@ class AlgorithmConfig: 'pretrained_model_name_or_path': MODELS_DIR['model2'], 'with_prior_preservation': True, 'prior_loss_weight': 1.0, - 'instance_prompt': 'a photo of sks person', - 'class_prompt': 'a photo of person', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of person', 'resolution': 512, 'train_batch_size': 1, 'gradient_accumulation_steps': 1, @@ -363,7 +363,7 @@ class AlgorithmConfig: 'seed': 0, 'mixed_precision': 'fp16', 'rank': 4, - 'validation_prompt': 'a photo of sks person', + 'validation_prompt': '((a selfie photo of person face)), distinct feature, head and shoulders shot, front view, face looking at camera', 'num_validation_images': 10, 'coords_log_interval': 10 } @@ -374,9 +374,9 @@ class AlgorithmConfig: 'conda_env': CONDA_ENVS['textual_inversion'], 'default_params': { 'pretrained_model_name_or_path': MODELS_DIR['model2'], - 'placeholder_token': '', + 'placeholder_token': '', 'initializer_token': 'person', - 'instance_prompt': 'a photo of person', + 'instance_prompt': 'a selfie photo of person', 'resolution': 512, 'train_batch_size': 1, 'gradient_accumulation_steps': 1, @@ -387,7 +387,7 @@ class AlgorithmConfig: 'checkpointing_steps': 500, 'seed': 0, 'mixed_precision': 'fp16', - 'validation_prompt': 'a close-up photo of person', + 'validation_prompt': 'a selfie photo of person', 'num_validation_images': 4, 'validation_epochs': 50, 'coords_log_interval': 10 diff --git a/src/backend/init_db.py b/src/backend/init_db.py index 1060fd6..055c1cb 100644 --- a/src/backend/init_db.py +++ b/src/backend/init_db.py @@ -93,18 +93,22 @@ def init_database(): data_types = [ { 'data_type_code': 'face', - 'instance_prompt': 'a photo of sks person', - 'class_prompt': 'a photo of person', - 'placeholder_token': '', + 'instance_prompt': 'a selfie photo of person', + 'class_prompt': 'a selfie photo of person', + 'validation_prompt_prefix_db_lora': '((a selfie photo of person face)), distinct feature, head and shoulders shot, front view, face looking at camera', + 'validation_prompt_prefix_ti': 'a selfie photo of person', + 'placeholder_token': '', 'initializer_token': 'person', 'description': '人脸类型的数据集' }, { 'data_type_code': 'art', - 'instance_prompt': 'a painting in style', + 'instance_prompt': 'a painting in style', 'class_prompt': 'a painting', - 'placeholder_token': '', - 'initializer_token': 'painting', + 'validation_prompt_prefix_db_lora': '((a painting in style)), distinct feature, coherent aesthetic, balanced lighting and color, fine art texture', + 'validation_prompt_prefix_ti': 'a painting in style', + 'placeholder_token': '', + 'initializer_token': 'style', 'description': '艺术品类型的数据集' } ]