fix: add submodule to aviod manully replace labml_nn dependencies in …

…docker

.dockerignore

@@ -19,3 +19,17 @@ package.sh
 compress.py
 
 logs/*
+
+shit.py
+
+load_m.py
+
+img2imgcopy.py
+
+img2img.py 
+
+img2img_test.py 
+
+gpuhold.py 
+
+compress.py  

.gitignore

@@ -7,12 +7,19 @@ log*/
 model_output/*
 outputs/
 model_output_test/
-
+tomesd/
 output_test/
 real_reg/*
 
 train_data/*
+train_data_crop/*
 
+our_image_outputs/
+our_imageoforigin2/
+our_json_imageoforigin2/
+our_json_outputs/
+other_models/*
+old_train_data/
 
 ## weights anf lib
 CLIP/
@@ -44,5 +51,51 @@ train_data_crop/*
 
 nnet.pt
 
+image_outputs/
+ImageReward/*
+
+home/
+
+final_test/
+
+final_json_data/
+
+final_evaluation/
+eval_prompts_advance/
+
+bench_samples/*
+
+bench_samples_standard/*
+
+aaaaaaasaveresults/
+
+abase_image_outputs/
+abase_json_outputs/
+
+test_data/
+
+test_json_out/
+
+combined_output.txt
+
+compress.__pycache__
+img2imgcopy.__pycache__
+load_m.__pycache__
+load_model_output.txt
+removebg.py 
+
+removeple_fnn.py
+run_local.sh 
+
+runDocker 
+sample_bench.py 
+score_local.py 
+
+score_oldone.py 
+score_output.txt 
+shit.py 
+shitshit.py 
+total_time.txt 
+
 
-other_models/stablediffusion/*
+gpuhold/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "dependency/annotated_deep_learning_paper_implementations"]
+	path = dependency/annotated_deep_learning_paper_implementations
+	url = https://github.com/SKDDJ/labml_nn.git

.vscode/launch.json

@@ -4,17 +4,63 @@
     // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
     "version": "0.2.0",
     "configurations": [
+
+
+
+
+        {
+            "name": "Python: score",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/score.py",
+            "console": "integratedTerminal",
+            "justMyCode": true
+        },
+        {
+            "name": "Python: loadmodel",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/load_model.py",
+            "args": [
+                // "--orig-img", "/home/schengwei/Competitionrepo/test_data/testdata/2_0.jpg",
+                // "--prompt", "a girl, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney,8k",
+                // "--scale", "0.99",
+                // "--steps", "25",
+                // "--strength", "0.03"
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": false
+        },
         {
-            "name": "Python: 当前文件",
+            "name": "Python: img2img",
             "type": "python",
             "request": "launch",
-            "program": "${file}",
-            "args": ["--hello", "hello world"],
+            "program": "${workspaceFolder}/img2img.py",
+            "args": [
+                "--orig-img", "/home/schengwei/Competitionrepo/test_data/testdata/2_0.jpg",
+                "--prompt", "a girl, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney,8k",
+                "--scale", "0.99",
+                "--steps", "25",
+                "--strength", "0.03"
+            ],
             "console": "integratedTerminal",
             "justMyCode": true
+        },
+        {
+            "name": "Python: main",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/main.py",
+            "args": [
+                // "--orig-img", "/home/schengwei/Competitionrepo/test_data/testdata/2_0.jpg",
+                // "--prompt", "a girl, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney,8k",
+                // "--scale", "0.99",
+                // "--steps", "25",
+                // "--strength", "0.03"
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": false
         }
-    ]
-}
 
-
-// accelerate launch --mixed_precision no --num_processes 1 trainn.py --instance_data_dir="train_data/newboy1"  --outdir="model_output/boy1"  --class_data_dir="real_reg/samples_boyface"   --with_prior_preservation  --prior_loss_weight=1.0  --class_prompt="boy" --num_class_images=200  --instance_prompt=" a <new1> boy"  --modifier_token "<new1>"
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "better-comments.highlightPlainText": true,
+    "files.exclude": {
+        "**/.idea": true
+    }
+}

Dockerfile

@@ -1,5 +1,6 @@
 # 使用基础镜像 pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
-FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
+# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
+FROM pytorch/pytorch:2.1.1-cuda12.1-cudnn8-runtime
 
 # 设置工作目录为 /workspace
 WORKDIR /workspace
@@ -15,7 +16,7 @@ COPY ./requirements.txt ./requirements.txt
 RUN apt update && apt upgrade -y
 
 # 安装所需的系统依赖包
-RUN apt install -y git libgl1-mesa-glx libglib2.0-0
+RUN DEBIAN_FRONTEND=noninteractive apt install -y git libgl1-mesa-glx libglib2.0-0
 
 # 安装gcc和g++编译器
 RUN apt-get install -y gcc g++
@@ -28,3 +29,5 @@ RUN pip install -r ./requirements.txt
 
 # 复制项目目录中的所有文件到容器的当前工作目录
 COPY . .
+
+

README.md

@@ -1,109 +1,41 @@
+
 # 代码说明
 
 ## 预测性能
 
-预计训练和推理总耗时10小时左右
+* **训练步数** ：预计训练 15,000 步
+* **耗时估计** ：总耗时约 5 小时
 
 ## 环境配置（必选）
 
-在 Dockerfile 中有每一步的注释，详情见 Dockerfile 文件
+* **Docker 配置** ：详细的步骤及注释在 `Dockerfile` 中。详见 `Dockerfile` 文件。
 
 ## 数据（必选）
 
-* 使用了CelebA（CelebFaces Attribute）人脸属性数据集，由香港中文大学开放提供，数据获取链接为http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html，随机获得200张男性人脸（在训练boy1和boy2模型时作为正则数据集使用）和200张女性人脸
-
-  * 在训练girl1和girl2模型时作为正则数据集使用
-* 对官方提供的数据集进行移除背景和裁剪人像的操作，放进train_data_crop文件夹内
-
-  * 在 train.sh文件中指定训练数据集的时候使用
+* **数据集** ：使用 Stable Diffusion webui 制作的高质量人像文本对数据集。
 
 ## 预训练模型（必选）
 
-* 使用了 CompVis/stable-diffusion-v1-4中的 AutoTokenizer， 代码继承自 custom diffusion中微调 CLIP Text Encoder 部分
-  * 路径位于other_models/stablediffusion/ 下
-  * 并且代码中应该会自动从 hugging face 上拉取
+* **模型** ：rembg
 
 ## 算法（必选）
 
-* 整体思路介绍（必选）
-
-[Texual Inversion]中微调CLIP Text Encoder 的方法 -> 为防止过拟合采用[DreamBooth]中的Prior Loss方法 -> 为 UViT 的指定 Layer 加入 LoRA
-
-* 方法的创新点
-
-之前在训练步数较少的情况下采用过 cross attention 的 LoRA
-
-* 算法的其他细节
-
-1. 自定义Token和CLIP Tokenizer更新：通过添加一个近义词的自定义token，扩展了CLIP模型的文本处理能力。更新CLIP Tokenizer，以包括新的token，确保模型能够识别和处理它。更新CLIPTextModel的embedding层，以适应新的token；
-2. 人脸正则数据集：为了避免模型的过拟合和语言漂移问题，引入了人脸正则数据集，帮助模型学习更好地理解人脸。训练数据集由 instance data 和 class data 组成，计算loss时，将模型输出分开，分别算loss；
-3. 参数微调：利用peft库中的lora，对模型的"qkv","fc1","fc2","proj","text_embed","clip_img_embed"部分进行参数微调。
+* **数据处理** ：使用 Stable Diffusion webui 制作的高质量人像文本对数据集来微调 unidiffuser 模型。
+* **算法增强** ：引入了 DDIM 算法以提高推理速度。
+* **图像处理** ：应用 image2image 方法对输入图像进行编辑。
 
 ### 整体思路介绍（必选）
 
-1. 用近义词初始化 `<new1>`，微调 CLIP Text Encoder
-2. 为了避免过拟合和语言漂移，加入人脸正则数据集；
-3. 利用peft库中的lora，对模型进行参数微调。
-
-训练代码说明:
-
-基于单个数据集进行训练:
-
-```shell
-accelerate launch --mixed_precision fp16 --num_processes 1 train.py -d '<训练文件所在位置>' -o '<模型输出>' --具体参数请参考 train.sh文件
-```
-
-训练所有数据集:
-
-```shell
-./train.sh
-```
-
-训练过程中的数据预处理, 模型架构, 加载方式等等均可进行修改, 只需要满足命令行接口即可, 并且模型输出位置的输出形式能够被 `sample.py`文件正确解析即可。
-
-生成代码说明:
-
-基于prompt文件生成图片:
-
-```shell
-python sample.py --restore_path '<模型输出>' --prompt_path '<prompt文件路径>' --output_path '<输出路径>'
-```
-
-基于所有的prompt文件进行生成:
-
-```
-./sample.sh
-```
-
-文件中sample的方式, prompt均可以更改, 但是测评时只会根据文件中的prompt进行测评。每个prompt要求输出三张图片, 并存放于目标路径。模型输出路径只需要能够正确解析训练产生的文件夹即可(若使用高效参数微调方法, 只需要将额外参数保存到输出路径并在sample.py中加载即可, 无需保存整个模型, 原模型可以从 `models/uvit_v1.pth`中加载)。`sample.py`除了会生成图片, 还会检查用于生成的模型和原模型的参数差异大小, 用于衡量微淘的参数量, 具体的计算方式见代码。
-
-### 客观指标打分(部分)
-
-```shell
-python score.py
-
-# 默认路径为
-    # --dataset: './train_data/'
-    # --prompts: './eval_prompts/'
-    # --outputs: './outputs/'
-
-# 可自行提供路径
-python score.py --dataset '<数据路径>' --prompts '<prompt路径>' --outputs '<提交路径>'
-```
+* **数据准备** ：使用 Stable Diffusion webui 制作高质量人像文本对数据集。
+* **速度优化** ：通过引入 DDIM 算法加快推理速度。
+* **功能实现** ：利用 image2image 方法实现对输入图像的编辑功能。
 
 ## 训练流程（必选）
 
-1. 直接在容器内运行. train.sh就可以开始训练
-   以下是代码运行逻辑
-2. 定义命令行参数，包括数据目录和输出目录。
-3. 在loop()函数中，使用train_state对象进行模型训练。
-4. 在训练过程中，使用accelerator.is_main_process判断当前进程是否为主进程。
-5. 如果是主进程，计算当前步数total_step，并在达到一定步数时记录日志和保存模型。
-6. 在达到一定步数时，保存模型的checkpoint文件，以便后续进行模型推理。
-7. 在训练结束时，保存最终的模型checkpoint文件。
-
-## 其他注意事项
-
-* 运行代码命令须参考修改过的train.sh和sample.sh文件
-* 控制训练步数（以“图文对”为单位）的参数所在位置：train.sh中的int 参数 train_step 来指定，默认 2000步
-* 因为我们初赛最终版本train_step为 10000步，并且使用了multi_stage的训练方法，即训练中途更换过训练集图片，但是这一版做了额外的优化，可能和初赛版本持平，若有问题请联系我们团队。
+1. **启动训练** ：在容器内运行 `train.sh` 脚本开始训练。
+2. **命令行参数** ：定义包括数据目录和输出目录在内的命令行参数。
+3. **训练逻辑** ：在 `loop()` 函数中利用 `train_state` 对象进行模型训练。
+4. **进程判断** ：使用 `accelerator.is_main_process` 判断是否为主进程。
+5. **日志记录** ：在主进程中，计算当前步数 `total_step`，并在达到一定步数时记录日志和保存模型。
+6. **模型保存** ：在达到指定步数时保存模型的 checkpoint 文件，以便后续进行模型推理。
+7. **训练结束** ：训练结束时保存最终的模型 checkpoint 文件。