mirror of
https://github.com/modelscope/modelscope.git
synced 2026-09-02 03:59:31 +02:00
Merge branch master-merge-github0829 into master
Title: merge github code Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13837366
This commit is contained in:
@@ -4,11 +4,15 @@ CODE_DIR=$PWD
|
||||
CODE_DIR_IN_CONTAINER=/Maas-lib
|
||||
echo "$USER"
|
||||
gpus='0,1 2,3 4,5 6,7'
|
||||
cpu_sets='45-58 31-44 16-30 0-15'
|
||||
cpu_sets='0-15 16-31 32-47 48-63'
|
||||
cpu_sets_arr=($cpu_sets)
|
||||
is_get_file_lock=false
|
||||
CI_COMMAND=${CI_COMMAND:-bash .dev_scripts/ci_container_test.sh python tests/run.py --parallel 2 --run_config tests/run_config.yaml}
|
||||
echo "ci command: $CI_COMMAND"
|
||||
PR_CHANGED_FILES="${PR_CHANGED_FILES:-''}"
|
||||
echo "PR modified files: $PR_CHANGED_FILES"
|
||||
PR_CHANGED_FILES=${PR_CHANGED_FILES//[ ]/#}
|
||||
echo "PR_CHANGED_FILES: $PR_CHANGED_FILES"
|
||||
idx=0
|
||||
for gpu in $gpus
|
||||
do
|
||||
@@ -42,6 +46,7 @@ do
|
||||
-e MODELSCOPE_ENVIRONMENT='ci' \
|
||||
-e TEST_UPLOAD_MS_TOKEN=$TEST_UPLOAD_MS_TOKEN \
|
||||
-e MODEL_TAG_URL=$MODEL_TAG_URL \
|
||||
-e PR_CHANGED_FILES=$PR_CHANGED_FILES \
|
||||
--workdir=$CODE_DIR_IN_CONTAINER \
|
||||
${IMAGE_NAME}:${IMAGE_VERSION} \
|
||||
$CI_COMMAND
|
||||
@@ -64,6 +69,7 @@ do
|
||||
-e MODELSCOPE_ENVIRONMENT='ci' \
|
||||
-e TEST_UPLOAD_MS_TOKEN=$TEST_UPLOAD_MS_TOKEN \
|
||||
-e MODEL_TAG_URL=$MODEL_TAG_URL \
|
||||
-e PR_CHANGED_FILES=$PR_CHANGED_FILES \
|
||||
--workdir=$CODE_DIR_IN_CONTAINER \
|
||||
${IMAGE_NAME}:${IMAGE_VERSION} \
|
||||
$CI_COMMAND
|
||||
|
||||
13
.github/workflows/citest.yaml
vendored
13
.github/workflows/citest.yaml
vendored
@@ -39,7 +39,7 @@ concurrency:
|
||||
jobs:
|
||||
unittest:
|
||||
# The type of runner that the job will run on
|
||||
runs-on: [modelscope-self-hosted]
|
||||
runs-on: [modelscope-self-hosted-us]
|
||||
timeout-minutes: 240
|
||||
steps:
|
||||
- name: ResetFileMode
|
||||
@@ -52,10 +52,19 @@ jobs:
|
||||
sudo chown -R $USER:$USER $ACTION_RUNNER_DIR
|
||||
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v2
|
||||
uses: actions/checkout@v3
|
||||
with:
|
||||
lfs: 'true'
|
||||
submodules: 'true'
|
||||
fetch-depth: ${{ github.event_name == 'pull_request' && 2 || 0 }}
|
||||
- name: Get changed files
|
||||
id: changed-files
|
||||
run: |
|
||||
if ${{ github.event_name == 'pull_request' }}; then
|
||||
echo "PR_CHANGED_FILES=$(git diff --name-only -r HEAD^1 HEAD | xargs)" >> $GITHUB_ENV
|
||||
else
|
||||
echo "PR_CHANGED_FILES=$(git diff --name-only ${{ github.event.before }} ${{ github.event.after }} | xargs)" >> $GITHUB_ENV
|
||||
fi
|
||||
- name: Checkout LFS objects
|
||||
run: git lfs checkout
|
||||
- name: Run unittest
|
||||
|
||||
2
.github/workflows/daily_regression.yaml
vendored
2
.github/workflows/daily_regression.yaml
vendored
@@ -12,7 +12,7 @@ concurrency:
|
||||
jobs:
|
||||
unittest:
|
||||
# The type of runner that the job will run on
|
||||
runs-on: [modelscope-self-hosted]
|
||||
runs-on: [modelscope-self-hosted-us]
|
||||
steps:
|
||||
- name: ResetFileMode
|
||||
shell: bash
|
||||
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -124,6 +124,8 @@ replace.sh
|
||||
result.png
|
||||
result.jpg
|
||||
result.mp4
|
||||
runs/
|
||||
ckpt/
|
||||
|
||||
# Pytorch
|
||||
*.pth
|
||||
|
||||
@@ -0,0 +1,456 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "04d4165c-fab2-4f54-9b50-11d53917d785",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# install required packages\n",
|
||||
"!pip install dashvector dashscope\n",
|
||||
"!pip install transformers_stream_generator python-dotenv"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "0ca135ac-b1b0-47b9-ad25-a0d11ac884f3",
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# prepare news corpus as knowledge source\n",
|
||||
"!git clone https://github.com/shijiebei2009/CEC-Corpus.git"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "728a2bf5-905c-48ef-b70a-be53d4f8fcc0",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:32:15.429699Z",
|
||||
"iopub.status.busy": "2023-08-10T10:32:15.429291Z",
|
||||
"iopub.status.idle": "2023-08-10T10:32:16.076518Z",
|
||||
"shell.execute_reply": "2023-08-10T10:32:16.075949Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:32:15.429679Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import dashscope\n",
|
||||
"import os\n",
|
||||
"from dotenv import load_dotenv\n",
|
||||
"from dashscope import TextEmbedding\n",
|
||||
"from dashvector import Client, Doc\n",
|
||||
"\n",
|
||||
"# get env variable from .env\n",
|
||||
"# please make sure DASHSCOPE_KEY is defined in .env\n",
|
||||
"load_dotenv()\n",
|
||||
"dashscope.api_key = os.getenv('DASHSCOPE_KEY')\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# initialize DashVector for embedding's indexing and searching\n",
|
||||
"dashvector_client = Client(api_key='{your-dashvector-api-key}')\n",
|
||||
"\n",
|
||||
"# define collection name\n",
|
||||
"collection_name = 'news_embeddings'\n",
|
||||
"\n",
|
||||
"# delete if already exist\n",
|
||||
"dashvector_client.delete(collection_name)\n",
|
||||
"\n",
|
||||
"# create a collection with embedding size of 1536\n",
|
||||
"rsp = dashvector_client.create(collection_name, 1536)\n",
|
||||
"collection = dashvector_client.get(collection_name)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"id": "558b64ab-1fdf-4339-8368-97e67bef8159",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:57:43.451192Z",
|
||||
"iopub.status.busy": "2023-08-10T10:57:43.450893Z",
|
||||
"iopub.status.idle": "2023-08-10T10:57:43.454858Z",
|
||||
"shell.execute_reply": "2023-08-10T10:57:43.454244Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:57:43.451173Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def prepare_data_from_dir(path, size):\n",
|
||||
" # prepare the data from a file folder in order to upsert to DashVector with a reasonable doc's size.\n",
|
||||
" batch_docs = []\n",
|
||||
" for file in os.listdir(path):\n",
|
||||
" with open(path + '/' + file, 'r', encoding='utf-8') as f:\n",
|
||||
" batch_docs.append(f.read())\n",
|
||||
" if len(batch_docs) == size:\n",
|
||||
" yield batch_docs[:]\n",
|
||||
" batch_docs.clear()\n",
|
||||
"\n",
|
||||
" if batch_docs:\n",
|
||||
" yield batch_docs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"id": "d65c0f3f-a080-4803-b5ed-f4e641a96db2",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:57:44.615001Z",
|
||||
"iopub.status.busy": "2023-08-10T10:57:44.614690Z",
|
||||
"iopub.status.idle": "2023-08-10T10:57:44.618899Z",
|
||||
"shell.execute_reply": "2023-08-10T10:57:44.618418Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:57:44.614979Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def prepare_data_from_file(path, size):\n",
|
||||
" # prepare the data from file in order to upsert to DashVector with a reasonable doc's size.\n",
|
||||
" batch_docs = []\n",
|
||||
" chunk_size = 12\n",
|
||||
" with open(path, 'r', encoding='utf-8') as f:\n",
|
||||
" doc = ''\n",
|
||||
" count = 0\n",
|
||||
" for line in f:\n",
|
||||
" if count < chunk_size and line.strip() != '':\n",
|
||||
" doc += line\n",
|
||||
" count += 1\n",
|
||||
" if count == chunk_size:\n",
|
||||
" batch_docs.append(doc)\n",
|
||||
" if len(batch_docs) == size:\n",
|
||||
" yield batch_docs[:]\n",
|
||||
" batch_docs.clear()\n",
|
||||
" doc = ''\n",
|
||||
" count = 0\n",
|
||||
"\n",
|
||||
" if batch_docs:\n",
|
||||
" yield batch_docs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"id": "aded6eec-1f05-479e-9f0e-3ce63872a07b",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:57:46.210192Z",
|
||||
"iopub.status.busy": "2023-08-10T10:57:46.209870Z",
|
||||
"iopub.status.idle": "2023-08-10T10:57:46.214412Z",
|
||||
"shell.execute_reply": "2023-08-10T10:57:46.213625Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:57:46.210172Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def generate_embeddings(docs):\n",
|
||||
" # create embeddings via DashScope's TextEmbedding model API\n",
|
||||
" rsp = TextEmbedding.call(model=TextEmbedding.Models.text_embedding_v1,\n",
|
||||
" input=docs)\n",
|
||||
" embeddings = [record['embedding'] for record in rsp.output['embeddings']]\n",
|
||||
" return embeddings if isinstance(docs, list) else embeddings[0]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5c0ba7e1-001f-4bb9-9bdb-7eb318bc3550",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"id = 0\n",
|
||||
"dir_name = 'CEC-Corpus/raw corpus/allSourceText'\n",
|
||||
"\n",
|
||||
"# indexing the raw docs with index to DashVector\n",
|
||||
"collection = dashvector_client.get(collection_name)\n",
|
||||
"\n",
|
||||
"# embedding api max batch size\n",
|
||||
"batch_size = 4 \n",
|
||||
"\n",
|
||||
"for news in list(prepare_data_from_dir(dir_name, batch_size)):\n",
|
||||
" ids = [id + i for i, _ in enumerate(news)]\n",
|
||||
" id += len(news)\n",
|
||||
" # generate embedding from raw docs\n",
|
||||
" vectors = generate_embeddings(news)\n",
|
||||
" # upsert and index\n",
|
||||
" ret = collection.upsert(\n",
|
||||
" [\n",
|
||||
" Doc(id=str(id), vector=vector, fields={\"raw\": doc})\n",
|
||||
" for id, doc, vector in zip(ids, news, vectors)\n",
|
||||
" ]\n",
|
||||
" )\n",
|
||||
" print(ret)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "53bed7e4-35be-4df6-8775-7d62fcdb6457",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# check the collection status\n",
|
||||
"collection = dashvector_client.get(collection_name)\n",
|
||||
"rsp = collection.stats()\n",
|
||||
"print(rsp)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"id": "41e54ddd-145d-49c3-ade4-4a46dc34e07b",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:57:54.368540Z",
|
||||
"iopub.status.busy": "2023-08-10T10:57:54.368215Z",
|
||||
"iopub.status.idle": "2023-08-10T10:57:54.371879Z",
|
||||
"shell.execute_reply": "2023-08-10T10:57:54.371364Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:57:54.368521Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def search_relevant_context(question, topk=1, client=dashvector_client):\n",
|
||||
" # query and recall the relevant information\n",
|
||||
" collection = client.get(collection_name)\n",
|
||||
"\n",
|
||||
" # recall the top k similarity results from DashVector\n",
|
||||
" rsp = collection.query(generate_embeddings(question), output_fields=['raw'],\n",
|
||||
" topk=topk)\n",
|
||||
" return \"\".join([item.fields['raw'] for item in rsp.output])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"id": "409236b9-87d4-4df0-8ee6-486d3c0e5fb6",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:57:56.141848Z",
|
||||
"iopub.status.busy": "2023-08-10T10:57:56.141502Z",
|
||||
"iopub.status.idle": "2023-08-10T10:57:56.387965Z",
|
||||
"shell.execute_reply": "2023-08-10T10:57:56.387379Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:57:56.141830Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"2006-08-26 10:41:45\n",
|
||||
"8月23日上午9时40分,京沪高速公路沧州服务区附近,一辆由北向南行驶的金杯面包车撞到高速公路护栏上,车上5名清华大学博士后研究人员及1名司机受伤,被紧急送往沧州二医院抢救。截至发稿时,仍有一名张姓博士后研究人员尚未脱离危险。\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# query the top 1 results\n",
|
||||
"question = '清华博士发生了什么?'\n",
|
||||
"context = search_relevant_context(question, topk=1)\n",
|
||||
"print(context)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "730abebb-1f5a-4fb9-b035-fb2ae09a31c9",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# initialize qwen 7B model\n",
|
||||
"from modelscope import AutoModelForCausalLM, AutoTokenizer\n",
|
||||
"from modelscope import GenerationConfig\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"qwen/Qwen-7B-Chat\", revision = 'v1.0.5',trust_remote_code=True)\n",
|
||||
"model = AutoModelForCausalLM.from_pretrained(\"qwen/Qwen-7B-Chat\", revision = 'v1.0.5',device_map=\"auto\", trust_remote_code=True, fp16=True).eval()\n",
|
||||
"model.generation_config = GenerationConfig.from_pretrained(\"Qwen/Qwen-7B-Chat\",revision = 'v1.0.5', trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "2f5a1bcb-e83a-44d3-bbe4-f97437782a3b",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:41:01.761863Z",
|
||||
"iopub.status.busy": "2023-08-10T10:41:01.761502Z",
|
||||
"iopub.status.idle": "2023-08-10T10:41:01.765849Z",
|
||||
"shell.execute_reply": "2023-08-10T10:41:01.765318Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:41:01.761842Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# define a prompt template for the vectorDB-enhanced LLM generation\n",
|
||||
"def answer_question(question, context):\n",
|
||||
" prompt = f'''请基于```内的内容回答问题。\"\n",
|
||||
"\t```\n",
|
||||
"\t{context}\n",
|
||||
"\t```\n",
|
||||
"\t我的问题是:{question}。\n",
|
||||
" '''\n",
|
||||
" history = None\n",
|
||||
" print(prompt)\n",
|
||||
" response, history = model.chat(tokenizer, prompt, history=None)\n",
|
||||
" return response"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "75ac8f4a-a861-4376-9e55-ebefef9a9cd6",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:41:29.070090Z",
|
||||
"iopub.status.busy": "2023-08-10T10:41:29.069778Z",
|
||||
"iopub.status.idle": "2023-08-10T10:41:31.613198Z",
|
||||
"shell.execute_reply": "2023-08-10T10:41:31.612421Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:41:29.070073Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"请基于```内的内容回答问题。\"\n",
|
||||
"\t```\n",
|
||||
"\t\n",
|
||||
"\t```\n",
|
||||
"\t我的问题是:清华博士发生了什么?。\n",
|
||||
" \n",
|
||||
"question: 清华博士发生了什么?\n",
|
||||
"answer: 清华博士是指清华大学的博士研究生。作为一名AI语言模型,我无法获取个人的身份信息或具体事件,因此无法回答清华博士发生了什么。如果您需要了解更多相关信息,建议您查询相关媒体或官方网站。\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# test the case on plain LLM without vectorDB enhancement\n",
|
||||
"question = '清华博士发生了什么?'\n",
|
||||
"answer = answer_question(question, '')\n",
|
||||
"print(f'question: {question}\\n' f'answer: {answer}')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"id": "eca328fc-cd69-4e12-8448-f426f3314414",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T10:41:34.268896Z",
|
||||
"iopub.status.busy": "2023-08-10T10:41:34.268585Z",
|
||||
"iopub.status.idle": "2023-08-10T10:41:37.750128Z",
|
||||
"shell.execute_reply": "2023-08-10T10:41:37.749414Z",
|
||||
"shell.execute_reply.started": "2023-08-10T10:41:34.268878Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"请基于```内的内容回答问题。\"\n",
|
||||
"\t```\n",
|
||||
"\t2006-08-26 10:41:45\n",
|
||||
"8月23日上午9时40分,京沪高速公路沧州服务区附近,一辆由北向南行驶的金杯面包车撞到高速公路护栏上,车上5名清华大学博士后研究人员及1名司机受伤,被紧急送往沧州二医院抢救。截至发稿时,仍有一名张姓博士后研究人员尚未脱离危险。\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\t```\n",
|
||||
"\t我的问题是:清华博士发生了什么?。\n",
|
||||
" \n",
|
||||
"question: 清华博士发生了什么?\n",
|
||||
"answer: 8月23日上午9时40分,一辆由北向南行驶的金杯面包车撞到高速公路护栏上,车上5名清华大学博士后研究人员及1名司机受伤,被紧急送往沧州二医院抢救。\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# test the case with knowledge\n",
|
||||
"context = search_relevant_context(question, topk=1)\n",
|
||||
"answer = answer_question(question, context)\n",
|
||||
"print(f'question: {question}\\n' f'answer: {answer}')"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.8.16"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9678e0bc-97cd-45bc-bd38-8d79c6789325",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# install required packages\n",
|
||||
"!pip install langchain\n",
|
||||
"!pip install unstructured\n",
|
||||
"!pip install transformers_stream_generator"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "36410a7c-a334-4ba2-abde-1679ac938a2a",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from typing import List, Optional\n",
|
||||
"from langchain.llms.base import LLM\n",
|
||||
"from modelscope import AutoModelForCausalLM, AutoTokenizer\n",
|
||||
"from modelscope import GenerationConfig\n",
|
||||
"\n",
|
||||
"# initialize qwen 7B model\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"qwen/Qwen-7B-Chat\", revision = 'v1.0.5',trust_remote_code=True)\n",
|
||||
"model = AutoModelForCausalLM.from_pretrained(\"qwen/Qwen-7B-Chat\", revision = 'v1.0.5',device_map=\"auto\", trust_remote_code=True, fp16=True).eval()\n",
|
||||
"model.generation_config = GenerationConfig.from_pretrained(\"Qwen/Qwen-7B-Chat\",revision = 'v1.0.5', trust_remote_code=True) \n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# torch garbage collection\n",
|
||||
"def torch_gc():\n",
|
||||
" os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
|
||||
" DEVICE = \"cuda\"\n",
|
||||
" DEVICE_ID = \"0\"\n",
|
||||
" CUDA_DEVICE = f\"{DEVICE}:{DEVICE_ID}\" if DEVICE_ID else DEVICE\n",
|
||||
" a = torch.Tensor([1, 2])\n",
|
||||
" a = a.cuda()\n",
|
||||
" print(a)\n",
|
||||
"\n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" with torch.cuda.device(CUDA_DEVICE):\n",
|
||||
" torch.cuda.empty_cache()\n",
|
||||
" torch.cuda.ipc_collect()\n",
|
||||
"\n",
|
||||
"# wrap the qwen model with langchain LLM base class\n",
|
||||
"class QianWenChatLLM(LLM):\n",
|
||||
" max_length = 10000\n",
|
||||
" temperature: float = 0.01\n",
|
||||
" top_p = 0.9\n",
|
||||
"\n",
|
||||
" def __init__(self):\n",
|
||||
" super().__init__()\n",
|
||||
"\n",
|
||||
" @property\n",
|
||||
" def _llm_type(self):\n",
|
||||
" return \"ChatLLM\"\n",
|
||||
"\n",
|
||||
" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:\n",
|
||||
" print(prompt)\n",
|
||||
" response, history = model.chat(tokenizer, prompt, history=None)\n",
|
||||
" torch_gc()\n",
|
||||
" return response\n",
|
||||
" \n",
|
||||
"# create the qwen llm\n",
|
||||
"qwllm = QianWenChatLLM()\n",
|
||||
"print('@@@ qianwen LLM created')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ce46aa8d-d772-4990-b748-12872fac2473",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-11T03:49:17.451327Z",
|
||||
"iopub.status.busy": "2023-08-11T03:49:17.450867Z",
|
||||
"iopub.status.idle": "2023-08-11T03:49:18.960037Z",
|
||||
"shell.execute_reply": "2023-08-11T03:49:18.959128Z",
|
||||
"shell.execute_reply.started": "2023-08-11T03:49:17.451304Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import re\n",
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"from typing import Any, List\n",
|
||||
"from pydantic import BaseModel, Extra\n",
|
||||
"from langchain.chains import RetrievalQA\n",
|
||||
"from langchain.document_loaders import UnstructuredFileLoader,TextLoader\n",
|
||||
"from langchain.embeddings.base import Embeddings\n",
|
||||
"from langchain.prompts import PromptTemplate\n",
|
||||
"from langchain.text_splitter import CharacterTextSplitter\n",
|
||||
"from langchain.vectorstores import FAISS\n",
|
||||
"\n",
|
||||
"# define chinese text split logic for divided docs into reasonable size\n",
|
||||
"class ChineseTextSplitter(CharacterTextSplitter):\n",
|
||||
" def __init__(self, pdf: bool = False, sentence_size: int = 100, **kwargs):\n",
|
||||
" super().__init__(**kwargs)\n",
|
||||
" self.pdf = pdf\n",
|
||||
" self.sentence_size = sentence_size\n",
|
||||
"\n",
|
||||
" def split_text(self, text: str) -> List[str]: \n",
|
||||
" if self.pdf:\n",
|
||||
" text = re.sub(r\"\\n{3,}\", r\"\\n\", text)\n",
|
||||
" text = re.sub('\\s', \" \", text)\n",
|
||||
" text = re.sub(\"\\n\\n\", \"\", text)\n",
|
||||
"\n",
|
||||
" text = re.sub(r'([;;.!?。!?\\?])([^”’])', r\"\\1\\n\\2\", text) # 单字符断句符\n",
|
||||
" text = re.sub(r'(\\.{6})([^\"’”」』])', r\"\\1\\n\\2\", text) # 英文省略号\n",
|
||||
" text = re.sub(r'(\\…{2})([^\"’”」』])', r\"\\1\\n\\2\", text) # 中文省略号\n",
|
||||
" text = re.sub(r'([;;!?。!?\\?][\"’”」』]{0,2})([^;;!?,。!?\\?])', r'\\1\\n\\2', text)\n",
|
||||
" # 如果双引号前有终止符,那么双引号才是句子的终点,把分句符\\n放到双引号后,注意前面的几句都小心保留了双引号\n",
|
||||
" text = text.rstrip() # 段尾如果有多余的\\n就去掉它\n",
|
||||
" # 很多规则中会考虑分号;,但是这里我把它忽略不计,破折号、英文双引号等同样忽略,需要的再做些简单调整即可。\n",
|
||||
" ls = [i for i in text.split(\"\\n\") if i]\n",
|
||||
" for ele in ls:\n",
|
||||
" if len(ele) > self.sentence_size:\n",
|
||||
" ele1 = re.sub(r'([,,.][\"’”」』]{0,2})([^,,.])', r'\\1\\n\\2', ele)\n",
|
||||
" ele1_ls = ele1.split(\"\\n\")\n",
|
||||
" for ele_ele1 in ele1_ls:\n",
|
||||
" if len(ele_ele1) > self.sentence_size:\n",
|
||||
" ele_ele2 = re.sub(r'([\\n]{1,}| {2,}[\"’”」』]{0,2})([^\\s])', r'\\1\\n\\2', ele_ele1)\n",
|
||||
" ele2_ls = ele_ele2.split(\"\\n\")\n",
|
||||
" for ele_ele2 in ele2_ls:\n",
|
||||
" if len(ele_ele2) > self.sentence_size:\n",
|
||||
" ele_ele3 = re.sub('( [\"’”」』]{0,2})([^ ])', r'\\1\\n\\2', ele_ele2)\n",
|
||||
" ele2_id = ele2_ls.index(ele_ele2)\n",
|
||||
" ele2_ls = ele2_ls[:ele2_id] + [i for i in ele_ele3.split(\"\\n\") if i] + ele2_ls[\n",
|
||||
" ele2_id + 1:]\n",
|
||||
" ele_id = ele1_ls.index(ele_ele1)\n",
|
||||
" ele1_ls = ele1_ls[:ele_id] + [i for i in ele2_ls if i] + ele1_ls[ele_id + 1:]\n",
|
||||
"\n",
|
||||
" id = ls.index(ele)\n",
|
||||
" ls = ls[:id] + [i for i in ele1_ls if i] + ls[id + 1:]\n",
|
||||
" return ls\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# using modelscope text embedding method for embedding tool\n",
|
||||
"class ModelScopeEmbeddings(BaseModel, Embeddings):\n",
|
||||
" embed: Any\n",
|
||||
" model_id: str =\"damo/nlp_corom_sentence-embedding_english-base\"\n",
|
||||
" \"\"\"Model name to use.\"\"\"\n",
|
||||
"\n",
|
||||
" def __init__(self, **kwargs: Any):\n",
|
||||
" \"\"\"Initialize the modelscope\"\"\"\n",
|
||||
" super().__init__(**kwargs)\n",
|
||||
" try:\n",
|
||||
" from modelscope.models import Model\n",
|
||||
" from modelscope.pipelines import pipeline\n",
|
||||
" from modelscope.utils.constant import Tasks\n",
|
||||
" self.embed = pipeline(Tasks.sentence_embedding,model=self.model_id)\n",
|
||||
"\n",
|
||||
" except ImportError as e:\n",
|
||||
" raise ValueError(\n",
|
||||
" \"Could not import some python packages.\" \"Please install it with `pip install modelscope`.\"\n",
|
||||
" ) from e\n",
|
||||
"\n",
|
||||
" class Config:\n",
|
||||
" extra = Extra.forbid\n",
|
||||
"\n",
|
||||
" def embed_documents(self, texts: List[str]) -> List[List[float]]:\n",
|
||||
" texts = list(map(lambda x: x.replace(\"\\n\", \" \"), texts))\n",
|
||||
" inputs = {\"source_sentence\": texts}\n",
|
||||
" embeddings = self.embed(input=inputs)['text_embedding']\n",
|
||||
" return embeddings\n",
|
||||
"\n",
|
||||
" def embed_query(self, text: str) -> List[float]:\n",
|
||||
" text = text.replace(\"\\n\", \" \")\n",
|
||||
" inputs = {\"source_sentence\": [text]}\n",
|
||||
" embedding = self.embed(input=inputs)['text_embedding'][0]\n",
|
||||
" return embedding\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "ca3dc051-1b0b-4bec-b082-6e94b220a34d",
|
||||
"metadata": {
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-10T06:44:05.671065Z",
|
||||
"iopub.status.busy": "2023-08-10T06:44:05.670720Z",
|
||||
"iopub.status.idle": "2023-08-10T06:44:05.674188Z",
|
||||
"shell.execute_reply": "2023-08-10T06:44:05.673699Z",
|
||||
"shell.execute_reply.started": "2023-08-10T06:44:05.671045Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# define prompt template\n",
|
||||
"prompt_template = \"\"\"请基于```内的内容回答问题。\"\n",
|
||||
"\t```\n",
|
||||
"\t{context}\n",
|
||||
"\t```\n",
|
||||
"\t我的问题是:{question}。\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"prompt = PromptTemplate(template=prompt_template, input_variables=[\"context\", \"question\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "a41ff8b8-bf19-4766-8d90-af48c7dfda99",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": true
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# load the vector db and upsert docs with vector to db\n",
|
||||
"\n",
|
||||
"print('@@@ reading docs ...')\n",
|
||||
"sentence_size = 1600\n",
|
||||
"embeddings = ModelScopeEmbeddings(model_id=\"damo/nlp_corom_sentence-embedding_chinese-tiny\")\n",
|
||||
"\n",
|
||||
"filepath = \"../../../README_zh.md\"\n",
|
||||
"if filepath.lower().endswith(\".md\"):\n",
|
||||
" loader = UnstructuredFileLoader(filepath, mode=\"elements\")\n",
|
||||
" docs = loader.load()\n",
|
||||
"elif filepath.lower().endswith(\".txt\"):\n",
|
||||
" loader = TextLoader(filepath, autodetect_encoding=True)\n",
|
||||
" textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)\n",
|
||||
" docs = loader.load_and_split(textsplitter) \n",
|
||||
"\n",
|
||||
"db = FAISS.from_documents(docs, embeddings)\n",
|
||||
"print('@@@ reading doc done, vec db created.')\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# create knowledge chain\n",
|
||||
"kc = RetrievalQA.from_llm(llm=qwllm, retriever=db.as_retriever(search_kwargs={\"k\": 6}), prompt=prompt)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "c97b1a9e-6260-4429-8411-a3a2cddadb05",
|
||||
"metadata": {
|
||||
"ExecutionIndicator": {
|
||||
"show": false
|
||||
},
|
||||
"execution": {
|
||||
"iopub.execute_input": "2023-08-06T06:14:23.817772Z",
|
||||
"iopub.status.busy": "2023-08-06T06:14:23.817192Z",
|
||||
"iopub.status.idle": "2023-08-06T06:14:27.775706Z",
|
||||
"shell.execute_reply": "2023-08-06T06:14:27.775194Z",
|
||||
"shell.execute_reply.started": "2023-08-06T06:14:23.817734Z"
|
||||
},
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"请基于```内的内容回答问题。\"\n",
|
||||
"\t```\n",
|
||||
"\tContext:\n",
|
||||
"ModelScope Library为模型贡献者提供了必要的分层API,以便将来自 CV、NLP、语音、多模态以及科学计算的模型集成到ModelScope生态系统中。所有这些不同模型的实现都以一种简单统一访问的方式进行封装,用户只需几行代码即可完成模型推理、微调和评估。同时,灵活的模块化设计使得在必要时也可以自定义模型训练推理过程中的不同组件。\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"ModelScope 是一个“模型即服务”(MaaS)平台,旨在汇集来自AI社区的最先进的机器学习模型,并简化在实际应用中使用AI模型的流程。ModelScope库使开发人员能够通过丰富的API设计执行推理、训练和评估,从而促进跨不同AI领域的最先进模型的统一体验。\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"除了包含各种模型的实现之外,ModelScope Library还支持与ModelScope后端服务进行必要的交互,特别是与Model-Hub和Dataset-Hub的交互。这种交互促进了模型和数据集的管理在后台无缝执行,包括模型数据集查询、版本控制、缓存管理等。\n",
|
||||
"\t```\n",
|
||||
"\t我的问题是:modelscope是什么?。\n",
|
||||
"\n",
|
||||
"tensor([1., 2.], device='cuda:0')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# test the knowledge chain\n",
|
||||
"query = 'modelscope是什么?'\n",
|
||||
"result = kc({\"query\": query})\n",
|
||||
"print(result)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.8.16"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -3,14 +3,13 @@ import sys
|
||||
import types
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from swift import LoRAConfig, Swift
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from modelscope import (EpochBasedTrainer, MsDataset, TorchModel, TrainingArgs,
|
||||
build_dataset_from_file, snapshot_download)
|
||||
from modelscope.metainfo import Trainers
|
||||
from modelscope.preprocessors import TextGenerationTransformersPreprocessor
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.trainers import build_trainer
|
||||
|
||||
DEFAULT_PAD_TOKEN = '[PAD]'
|
||||
@@ -205,12 +204,12 @@ preprocessor = TextGenerationTransformersPreprocessor(
|
||||
|
||||
if args.use_lora != 0:
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=['pack'],
|
||||
rank=args.lora_rank,
|
||||
target_modules=['pack'],
|
||||
r=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout)
|
||||
model = model.bfloat16()
|
||||
Swift.prepare_model(model, lora_config)
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
|
||||
kwargs = dict(
|
||||
model=model,
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.utils.constant import Tasks
|
||||
|
||||
# 使用源模型 model_id 初始化 pipeline
|
||||
@@ -12,11 +11,11 @@ model_id = 'baichuan-inc/baichuan-7B'
|
||||
pipe = pipeline(
|
||||
task=Tasks.text_generation, model=model_id, model_revision='v1.0.2')
|
||||
# lora 配置,replace_modules,rank,alpha 需与训练参数相同
|
||||
lora_config = LoRAConfig(replace_modules=['pack'], rank=32, lora_alpha=32)
|
||||
lora_config = LoRAConfig(target_modules=['pack'], r=32, lora_alpha=32)
|
||||
# 转 bf16,需与训练精度相同
|
||||
model = pipe.model.bfloat16()
|
||||
# model 转 lora
|
||||
Swift.prepare_model(model, lora_config)
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
# 加载 lora 参数,默认 link 到于 output/model 路径
|
||||
work_dir = './tmp'
|
||||
state_dict = torch.load(osp.join(work_dir, 'output/pytorch_model.bin'))
|
||||
|
||||
@@ -4,6 +4,7 @@ from dataclasses import dataclass, field
|
||||
import numpy as np
|
||||
import torch
|
||||
from chatglm_trainer import Seq2SeqTrainer
|
||||
from swift import LoRAConfig, Swift
|
||||
from text_generation_metric import TextGenerationMetric
|
||||
from transformers import DataCollatorForSeq2Seq
|
||||
|
||||
@@ -11,8 +12,6 @@ from modelscope import build_dataset_from_file, snapshot_download
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.models import Model
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.trainers.training_args import TrainingArgs
|
||||
from modelscope.utils.config import ConfigDict
|
||||
from modelscope.utils.hub import read_config
|
||||
@@ -243,15 +242,15 @@ elif not args.use_lora:
|
||||
|
||||
if args.use_lora != 0:
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=['attention.query_key_value'],
|
||||
rank=args.lora_rank,
|
||||
target_modules=['attention.query_key_value'],
|
||||
r=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout)
|
||||
if args.use_amp:
|
||||
model = model.float()
|
||||
else:
|
||||
model = model.bfloat16()
|
||||
Swift.prepare_model(model, lora_config)
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
|
||||
prefix = args.source_prefix if args.source_prefix is not None else ''
|
||||
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
|
||||
from modelscope import Model, pipeline, read_config
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.utils.config import ConfigDict
|
||||
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=['attention.query_key_value'],
|
||||
rank=32,
|
||||
target_modules=['attention.query_key_value'],
|
||||
r=32,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
pretrained_weights='./lora_dureader_target/iter_600.pth')
|
||||
lora_dropout=0.05)
|
||||
|
||||
model_dir = 'ZhipuAI/ChatGLM-6B'
|
||||
model_config = read_config(model_dir)
|
||||
@@ -19,8 +21,12 @@ model_config['model'] = ConfigDict({
|
||||
|
||||
model = Model.from_pretrained(model_dir, cfg_dict=model_config)
|
||||
model = model.bfloat16()
|
||||
Swift.prepare_model(model, lora_config)
|
||||
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
work_dir = './tmp'
|
||||
state_dict = torch.load(osp.join(work_dir, 'iter_600.pth'))
|
||||
model = Swift.from_pretrained(
|
||||
model, osp.join(work_dir, 'output_best'), device_map='auto')
|
||||
model.load_state_dict(state_dict)
|
||||
pipe = pipeline('chat', model, pipeline_name='chatglm6b-text-generation')
|
||||
|
||||
print(
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
|
||||
from modelscope import Model, pipeline, read_config
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.utils.config import ConfigDict
|
||||
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=['attention.query_key_value'],
|
||||
rank=32,
|
||||
target_modules=['attention.query_key_value'],
|
||||
r=32,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
pretrained_weights='./lora_dureader_target/iter_600.pth')
|
||||
lora_dropout=0.05)
|
||||
|
||||
model_dir = 'ZhipuAI/chatglm2-6b'
|
||||
model_config = read_config(model_dir)
|
||||
@@ -19,7 +21,12 @@ model_config['model'] = ConfigDict({
|
||||
|
||||
model = Model.from_pretrained(model_dir, cfg_dict=model_config)
|
||||
model = model.bfloat16()
|
||||
Swift.prepare_model(model, lora_config)
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
work_dir = './tmp'
|
||||
state_dict = torch.load(osp.join(work_dir, 'iter_600.pth'))
|
||||
model = Swift.from_pretrained(
|
||||
model, osp.join(work_dir, 'output_best'), device_map='auto')
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
pipe = pipeline('chat', model, pipeline_name='chatglm2_6b-text-generation')
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
|
||||
|
||||
import json
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
|
||||
from modelscope import TrainingArgs, build_dataset_from_file
|
||||
from modelscope.hub.snapshot_download import snapshot_download
|
||||
@@ -16,8 +17,6 @@ from modelscope.models.nlp.llama import LlamaForTextGeneration, LlamaTokenizer
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.msdatasets.dataset_cls.custom_datasets.torch_custom_dataset import \
|
||||
TorchCustomDataset
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.trainers import build_trainer
|
||||
|
||||
IGNORE_INDEX = -100
|
||||
@@ -335,12 +334,12 @@ if __name__ == '__main__':
|
||||
|
||||
if args.use_lora != 0:
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=['q_proj', 'k_proj', 'v_proj'],
|
||||
rank=args.lora_rank,
|
||||
target_modules=['q_proj', 'k_proj', 'v_proj'],
|
||||
r=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout)
|
||||
model = model.bfloat16()
|
||||
Swift.prepare_model(model, lora_config)
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
|
||||
tokenizer = LlamaTokenizer.from_pretrained(
|
||||
model_path,
|
||||
|
||||
81
examples/pytorch/llm/README.md
Normal file
81
examples/pytorch/llm/README.md
Normal file
@@ -0,0 +1,81 @@
|
||||
<h1 align="center">LLM SFT Example</h1>
|
||||
|
||||
<p align="center">
|
||||
<img src="https://img.shields.io/badge/python-%E2%89%A53.8-5be.svg">
|
||||
<img src="https://img.shields.io/badge/pytorch-%E2%89%A51.12%20%7C%20%E2%89%A52.0-orange.svg">
|
||||
<a href="https://github.com/modelscope/modelscope/"><img src="https://img.shields.io/badge/modelscope-%E2%89%A51.8.1-5D91D4.svg"></a>
|
||||
<a href="https://github.com/modelscope/swift/"><img src="https://img.shields.io/badge/ms--swift-%E2%89%A51.0.0-6FEBB9.svg"></a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://modelscope.cn/home">Modelscope Hub</a>
|
||||
<br>
|
||||
<a href="README_CN.md">中文</a>  |  English
|
||||
</p>
|
||||
|
||||
## Note
|
||||
1. This README.md file is **copied from** [ms-swift](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm/README.md)
|
||||
2. This directory has been **migrated** to [ms-swift](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm), and the files in this directory are **no longer maintained**.
|
||||
|
||||
## Features
|
||||
1. supported sft method: [lora](https://arxiv.org/abs/2106.09685), [qlora](https://arxiv.org/abs/2305.14314), full(full parameter fine tuning), ...
|
||||
2. supported models: [**qwen-7b**](https://github.com/QwenLM/Qwen-7B), baichuan-7b, baichuan-13b, chatglm2-6b, chatglm2-6b-32k, llama2-7b, llama2-13b, llama2-70b, openbuddy-llama2-13b, openbuddy-llama-65b, polylm-13b, ...
|
||||
3. supported feature: quantization, ddp, model parallelism(device map), gradient checkpoint, gradient accumulation steps, push to modelscope hub, custom datasets, ...
|
||||
4. supported datasets: alpaca-en(gpt4), alpaca-zh(gpt4), finance-en, multi-alpaca-all, code-en, instinwild-en, instinwild-zh, ...
|
||||
|
||||
## Prepare the Environment
|
||||
Experimental environment: A10, 3090, A100, ... (V100 does not support bf16, quantization)
|
||||
```bash
|
||||
# Installing miniconda
|
||||
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
|
||||
sh Miniconda3-latest-Linux-x86_64.sh
|
||||
|
||||
# Setting up a conda virtual environment
|
||||
conda create --name ms-sft python=3.10
|
||||
conda activate ms-sft
|
||||
|
||||
# Setting up a global pip mirror for faster downloads
|
||||
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
pip install torch torchvision torchaudio -U
|
||||
pip install sentencepiece charset_normalizer cpm_kernels tiktoken -U
|
||||
pip install matplotlib scikit-learn tqdm tensorboard -U
|
||||
pip install transformers datasets -U
|
||||
pip install accelerate transformers_stream_generator -U
|
||||
|
||||
pip install ms-swift modelscope -U
|
||||
# Recommended installation from source code for faster bug fixes
|
||||
git clone https://github.com/modelscope/swift.git
|
||||
cd swift
|
||||
pip install -r requirements.txt
|
||||
pip install .
|
||||
# same as modelscope...(git clone ...)
|
||||
```
|
||||
|
||||
## Run SFT and Inference
|
||||
```bash
|
||||
# Clone the repository and enter the code directory.
|
||||
git clone https://github.com/modelscope/swift.git
|
||||
cd swift/examples/pytorch/llm
|
||||
|
||||
# sft(qlora) and infer qwen-7b, Requires 16GB VRAM.
|
||||
# If you want to use quantification, you need to `pip install bitsandbytes`
|
||||
bash scripts/qwen_7b/qlora/sft.sh
|
||||
# If you want to push the model to modelscope hub during training
|
||||
bash scripts/qwen_7b/qlora/sft_push_to_hub.sh
|
||||
bash scripts/qwen_7b/qlora/infer.sh
|
||||
|
||||
# sft(qlora+ddp) and infer qwen-7b, Requires 4*16GB VRAM.
|
||||
bash scripts/qwen_7b/qlora_ddp/sft.sh
|
||||
bash scripts/qwen_7b/qlora_ddp/infer.sh
|
||||
|
||||
# sft(full) and infer qwen-7b, Requires 95GB VRAM.
|
||||
bash scripts/qwen_7b/full/sft.sh
|
||||
bash scripts/qwen_7b/full/infer.sh
|
||||
|
||||
# For more scripts, please see `scripts/` folder
|
||||
```
|
||||
|
||||
## Extend Datasets
|
||||
1. If you need to extend the model, you can modify the `MODEL_MAPPING` in `utils/models.py`. `model_id` can be specified as a local path. In this case, `revision` doesn't work.
|
||||
2. If you need to extend or customize the dataset, you can modify the `DATASET_MAPPING` in `utils/datasets.py`. You need to customize the `get_*_dataset` function, which returns a dataset with two columns: `instruction`, `output`.
|
||||
83
examples/pytorch/llm/README_CN.md
Normal file
83
examples/pytorch/llm/README_CN.md
Normal file
@@ -0,0 +1,83 @@
|
||||
<h1 align="center">大模型微调的例子</h1>
|
||||
|
||||
<p align="center">
|
||||
<img src="https://img.shields.io/badge/python-%E2%89%A53.8-5be.svg">
|
||||
<img src="https://img.shields.io/badge/pytorch-%E2%89%A51.12%20%7C%20%E2%89%A52.0-orange.svg">
|
||||
<a href="https://github.com/modelscope/modelscope/"><img src="https://img.shields.io/badge/modelscope-%E2%89%A51.8.1-5D91D4.svg"></a>
|
||||
<a href="https://github.com/modelscope/swift/"><img src="https://img.shields.io/badge/ms--swift-%E2%89%A51.0.0-6FEBB9.svg">
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://modelscope.cn/home">魔搭社区</a>
|
||||
<br>
|
||||
中文  |  <a href="README.md">English</a>
|
||||
</p>
|
||||
|
||||
## 请注意
|
||||
1. 该README_CN.md**拷贝**自[ms-swift](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm/README_CN.md)
|
||||
2. 该目录已经**迁移**至[ms-swift](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm), 此目录中的文件**不再维护**.
|
||||
|
||||
## 特性
|
||||
1. [lora](https://arxiv.org/abs/2106.09685), [qlora](https://arxiv.org/abs/2305.14314), 全参数微调, ...
|
||||
2. 支持的模型: [**qwen-7b**](https://github.com/QwenLM/Qwen-7B), baichuan-7b, baichuan-13b, chatglm2-6b, chatglm2-6b-32k, llama2-7b, llama2-13b, llama2-70b, openbuddy-llama2-13b, openbuddy-llama-65b, polylm-13b, ...
|
||||
3. 支持的特性: 模型量化, DDP, 模型并行(device_map), gradient checkpoint, 梯度累加, 支持推送modelscope hub, 支持自定义数据集, ...
|
||||
4. 支持的数据集: alpaca-en(gpt4), alpaca-zh(gpt4), finance-en, multi-alpaca-all, code-en, instinwild-en, instinwild-zh, ...
|
||||
|
||||
|
||||
## 准备实验环境
|
||||
实验环境: A10, 3090, A100均可. (V100不支持bf16, 量化)
|
||||
```bash
|
||||
# 安装miniconda
|
||||
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
|
||||
# 一直[ENTER], 最后一个选项yes即可
|
||||
sh Miniconda3-latest-Linux-x86_64.sh
|
||||
|
||||
# conda虚拟环境搭建
|
||||
conda create --name ms-sft python=3.10
|
||||
conda activate ms-sft
|
||||
|
||||
# pip设置全局镜像与相关python包安装
|
||||
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
pip install torch torchvision torchaudio -U
|
||||
pip install sentencepiece charset_normalizer cpm_kernels tiktoken -U
|
||||
pip install matplotlib scikit-learn tqdm tensorboard -U
|
||||
pip install transformers datasets -U
|
||||
pip install accelerate transformers_stream_generator -U
|
||||
|
||||
pip install ms-swift modelscope -U
|
||||
# 推荐从源码安装swift和modelscope, 这具有更多的特性和更快的bug修复
|
||||
git clone https://github.com/modelscope/swift.git
|
||||
cd swift
|
||||
pip install -r requirements.txt
|
||||
pip install .
|
||||
# modelscope类似...(git clone ...)
|
||||
```
|
||||
|
||||
## 微调和推理
|
||||
```bash
|
||||
# clone仓库并进入代码目录
|
||||
git clone https://github.com/modelscope/swift.git
|
||||
cd swift/examples/pytorch/llm
|
||||
|
||||
# 微调(qlora)+推理 qwen-7b, 需要16GB显存.
|
||||
# 如果你想要使用量化, 你需要`pip install bitsandbytes`
|
||||
bash scripts/qwen_7b/qlora/sft.sh
|
||||
# 如果你想在训练时, 将权重push到modelscope hub中.
|
||||
bash scripts/qwen_7b/qlora/sft_push_to_hub.sh
|
||||
bash scripts/qwen_7b/qlora/infer.sh
|
||||
|
||||
# 微调(qlora+ddp)+推理 qwen-7b, 需要4卡*16GB显存.
|
||||
bash scripts/qwen_7b/qlora_ddp/sft.sh
|
||||
bash scripts/qwen_7b/qlora_ddp/infer.sh
|
||||
|
||||
# 微调(full)+推理 qwen-7b, 需要95G显存.
|
||||
bash scripts/qwen_7b/full/sft.sh
|
||||
bash scripts/qwen_7b/full/infer.sh
|
||||
|
||||
# 更多的scripts脚本, 可以看`scripts`文件夹
|
||||
```
|
||||
|
||||
## 拓展数据集
|
||||
1. 如果你想要拓展模型, 你可以修改`utils/models.py`文件中的`MODEL_MAPPING`. `model_id`可以指定为本地路径, 这种情况下, `revision`参数不起作用.
|
||||
2. 如果你想要拓展或使用自定义数据集, 你可以修改`utils/datasets.py`文件中的`DATASET_MAPPING`. 你需要自定义`get_*_dataset`函数, 并返回包含`instruction`, `output`两列的数据集.
|
||||
@@ -1,18 +1,25 @@
|
||||
# ### Setting up experimental environment.
|
||||
import os
|
||||
# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
from transformers import GenerationConfig, TextStreamer
|
||||
from utils import (DATASET_MAPPER, DEFAULT_PROMPT, MODEL_MAPPER, get_dataset,
|
||||
from utils import (DATASET_MAPPING, DEFAULT_PROMPT, MODEL_MAPPING, get_dataset,
|
||||
get_model_tokenizer, inference, parse_args, process_dataset,
|
||||
tokenize_function)
|
||||
|
||||
from modelscope import get_logger
|
||||
from modelscope.swift import LoRAConfig, Swift
|
||||
|
||||
warnings.warn(
|
||||
'This directory has been migrated to '
|
||||
'https://github.com/modelscope/swift/tree/main/examples/pytorch/llm, '
|
||||
'and the files in this directory are no longer maintained.',
|
||||
DeprecationWarning)
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
@@ -20,7 +27,7 @@ logger = get_logger()
|
||||
@dataclass
|
||||
class InferArguments:
|
||||
model_type: str = field(
|
||||
default='qwen-7b', metadata={'choices': list(MODEL_MAPPER.keys())})
|
||||
default='qwen-7b', metadata={'choices': list(MODEL_MAPPING.keys())})
|
||||
sft_type: str = field(
|
||||
default='lora', metadata={'choices': ['lora', 'full']})
|
||||
ckpt_path: str = '/path/to/your/iter_xxx.pth'
|
||||
@@ -29,9 +36,9 @@ class InferArguments:
|
||||
|
||||
dataset: str = field(
|
||||
default='alpaca-en,alpaca-zh',
|
||||
metadata={'help': f'dataset choices: {list(DATASET_MAPPER.keys())}'})
|
||||
metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})
|
||||
dataset_seed: int = 42
|
||||
dataset_sample: Optional[int] = None
|
||||
dataset_sample: int = 20000 # -1: all dataset
|
||||
dataset_test_size: float = 0.01
|
||||
prompt: str = DEFAULT_PROMPT
|
||||
max_length: Optional[int] = 2048
|
||||
@@ -48,7 +55,8 @@ class InferArguments:
|
||||
|
||||
def __post_init__(self):
|
||||
if self.lora_target_modules is None:
|
||||
self.lora_target_modules = MODEL_MAPPER[self.model_type]['lora_TM']
|
||||
self.lora_target_modules = MODEL_MAPPING[
|
||||
self.model_type]['lora_TM']
|
||||
|
||||
if not os.path.isfile(self.ckpt_path):
|
||||
raise ValueError(
|
||||
@@ -60,19 +68,23 @@ def llm_infer(args: InferArguments) -> None:
|
||||
support_bf16 = torch.cuda.is_bf16_supported()
|
||||
if not support_bf16:
|
||||
logger.warning(f'support_bf16: {support_bf16}')
|
||||
|
||||
kwargs = {'low_cpu_mem_usage': True, 'device_map': 'auto'}
|
||||
model, tokenizer, _ = get_model_tokenizer(
|
||||
args.model_type, torch_dtype=torch.bfloat16)
|
||||
args.model_type, torch_dtype=torch.bfloat16, **kwargs)
|
||||
|
||||
# ### Preparing lora
|
||||
if args.sft_type == 'lora':
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=args.lora_target_modules,
|
||||
rank=args.lora_rank,
|
||||
target_modules=args.lora_target_modules,
|
||||
r=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout_p,
|
||||
pretrained_weights=args.ckpt_path)
|
||||
logger.info(f'lora_config: {lora_config}')
|
||||
model = Swift.prepare_model(model, lora_config)
|
||||
state_dict = torch.load(args.ckpt_path, map_location='cpu')
|
||||
model.load_state_dict(state_dict)
|
||||
elif args.sft_type == 'full':
|
||||
state_dict = torch.load(args.ckpt_path, map_location='cpu')
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
@@ -2,9 +2,8 @@
|
||||
"""
|
||||
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y
|
||||
pip install sentencepiece charset_normalizer cpm_kernels tiktoken -U
|
||||
pip install matplotlib scikit-learn -U
|
||||
pip install transformers datasets -U
|
||||
pip install tqdm tensorboard torchmetrics -U
|
||||
pip install transformers datasets scikit-learn -U
|
||||
pip install matplotlib tqdm tensorboard torchmetrics -U
|
||||
pip install accelerate transformers_stream_generator -U
|
||||
|
||||
# Install the latest version of modelscope from source
|
||||
@@ -14,14 +13,16 @@ pip install -r requirements.txt
|
||||
pip install .
|
||||
"""
|
||||
import os
|
||||
# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from types import MethodType
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from swift import LoRAConfig, Swift
|
||||
from torch import Tensor
|
||||
from utils import (DATASET_MAPPER, DEFAULT_PROMPT, MODEL_MAPPER,
|
||||
from utils import (DATASET_MAPPING, DEFAULT_PROMPT, MODEL_MAPPING,
|
||||
data_collate_fn, get_dataset, get_model_tokenizer,
|
||||
get_T_max, get_work_dir, parse_args, plot_images,
|
||||
print_example, print_model_info, process_dataset,
|
||||
@@ -29,10 +30,14 @@ from utils import (DATASET_MAPPER, DEFAULT_PROMPT, MODEL_MAPPER,
|
||||
tokenize_function)
|
||||
|
||||
from modelscope import get_logger
|
||||
from modelscope.swift import LoRAConfig, Swift
|
||||
from modelscope.trainers import EpochBasedTrainer
|
||||
from modelscope.utils.config import Config
|
||||
|
||||
warnings.warn(
|
||||
'This directory has been migrated to '
|
||||
'https://github.com/modelscope/swift/tree/main/examples/pytorch/llm, '
|
||||
'and the files in this directory are no longer maintained.',
|
||||
DeprecationWarning)
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@@ -40,7 +45,7 @@ logger = get_logger()
|
||||
class SftArguments:
|
||||
seed: int = 42
|
||||
model_type: str = field(
|
||||
default='qwen-7b', metadata={'choices': list(MODEL_MAPPER.keys())})
|
||||
default='qwen-7b', metadata={'choices': list(MODEL_MAPPING.keys())})
|
||||
# baichuan-7b: 'lora': 16G; 'full': 80G
|
||||
sft_type: str = field(
|
||||
default='lora', metadata={'choices': ['lora', 'full']})
|
||||
@@ -49,9 +54,9 @@ class SftArguments:
|
||||
|
||||
dataset: str = field(
|
||||
default='alpaca-en,alpaca-zh',
|
||||
metadata={'help': f'dataset choices: {list(DATASET_MAPPER.keys())}'})
|
||||
metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})
|
||||
dataset_seed: int = 42
|
||||
dataset_sample: Optional[int] = None
|
||||
dataset_sample: int = 20000 # -1: all dataset
|
||||
dataset_test_size: float = 0.01
|
||||
prompt: str = DEFAULT_PROMPT
|
||||
max_length: Optional[int] = 2048
|
||||
@@ -78,6 +83,13 @@ class SftArguments:
|
||||
logging_interval: int = 5
|
||||
tb_interval: int = 5
|
||||
|
||||
# other
|
||||
use_flash_attn: Optional[bool] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
'help': "This parameter is used only when model_type='qwen-7b'"
|
||||
})
|
||||
|
||||
def __post_init__(self):
|
||||
if self.sft_type == 'lora':
|
||||
if self.learning_rate is None:
|
||||
@@ -102,7 +114,10 @@ class SftArguments:
|
||||
self.output_dir = os.path.join(self.output_dir, self.model_type)
|
||||
|
||||
if self.lora_target_modules is None:
|
||||
self.lora_target_modules = MODEL_MAPPER[self.model_type]['lora_TM']
|
||||
self.lora_target_modules = MODEL_MAPPING[
|
||||
self.model_type]['lora_TM']
|
||||
if self.use_flash_attn is None:
|
||||
self.use_flash_attn = 'auto'
|
||||
|
||||
|
||||
def llm_sft(args: SftArguments) -> None:
|
||||
@@ -112,22 +127,22 @@ def llm_sft(args: SftArguments) -> None:
|
||||
support_bf16 = torch.cuda.is_bf16_supported()
|
||||
if not support_bf16:
|
||||
logger.warning(f'support_bf16: {support_bf16}')
|
||||
|
||||
kwargs = {'low_cpu_mem_usage': True, 'device_map': 'auto'}
|
||||
if args.model_type == 'qwen-7b':
|
||||
kwargs['use_flash_attn'] = args.use_flash_attn
|
||||
model, tokenizer, model_dir = get_model_tokenizer(
|
||||
args.model_type, torch_dtype=torch.bfloat16)
|
||||
args.model_type, torch_dtype=torch.bfloat16, **kwargs)
|
||||
|
||||
if args.gradient_checkpoint:
|
||||
# baichuan-13b does not implement the `get_input_embeddings` function
|
||||
if args.model_type == 'baichuan-13b':
|
||||
model.get_input_embeddings = MethodType(
|
||||
lambda self: self.model.embed_tokens, model)
|
||||
model.gradient_checkpointing_enable()
|
||||
model.enable_input_require_grads()
|
||||
|
||||
# ### Preparing lora
|
||||
if args.sft_type == 'lora':
|
||||
lora_config = LoRAConfig(
|
||||
replace_modules=args.lora_target_modules,
|
||||
rank=args.lora_rank,
|
||||
target_modules=args.lora_target_modules,
|
||||
r=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout_p)
|
||||
logger.info(f'lora_config: {lora_config}')
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from .dataset import DATASET_MAPPER, get_dataset, process_dataset
|
||||
from .models import MODEL_MAPPER, get_model_tokenizer
|
||||
from .dataset import DATASET_MAPPING, get_dataset, process_dataset
|
||||
from .models import MODEL_MAPPING, get_model_tokenizer
|
||||
from .utils import (DEFAULT_PROMPT, MyMetric, data_collate_fn, get_T_max,
|
||||
get_work_dir, inference, parse_args, plot_images,
|
||||
print_example, print_model_info, read_tensorboard_file,
|
||||
|
||||
@@ -1,107 +1,146 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from functools import partial
|
||||
from typing import Callable, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from datasets import Dataset as HfDataset
|
||||
from datasets import concatenate_datasets
|
||||
from numpy.random import RandomState
|
||||
from swift.utils import get_seed
|
||||
|
||||
from modelscope import MsDataset
|
||||
|
||||
|
||||
def _processing_alpaca(dataset: HfDataset) -> HfDataset:
|
||||
def _processing_alpaca(
|
||||
dataset: HfDataset,
|
||||
preprocess_input: Optional[Callable[[str], str]] = None) -> HfDataset:
|
||||
instruction = dataset['instruction']
|
||||
input_ = dataset['input']
|
||||
res = []
|
||||
new_instruction = []
|
||||
for inst, inp in zip(instruction, input_):
|
||||
if inp is not None and inp != '':
|
||||
if inp.startswith('输入:'):
|
||||
inp = inp[3:]
|
||||
inst = f'{inst}\n{inp}'
|
||||
res.append(inst)
|
||||
if inp is None:
|
||||
inp = ''
|
||||
if preprocess_input is not None:
|
||||
inp = preprocess_input(inp)
|
||||
inst = f'{inst}\n{inp}'
|
||||
new_instruction.append(inst)
|
||||
dataset = HfDataset.from_dict({
|
||||
'instruction': res,
|
||||
'instruction': new_instruction,
|
||||
'output': dataset['output']
|
||||
})
|
||||
return dataset
|
||||
|
||||
|
||||
def _processing_multi_alpaca(datasets: [HfDataset, List]) -> HfDataset:
|
||||
output = []
|
||||
res = []
|
||||
def get_alpaca_gpt4_en_dataset() -> HfDataset:
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'AI-ModelScope/alpaca-gpt4-data-en', split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
if not isinstance(datasets, List):
|
||||
datasets = [datasets]
|
||||
for dataset in datasets:
|
||||
instruction = dataset['instruction']
|
||||
input_ = dataset['input']
|
||||
output_ = dataset['output']
|
||||
for inst, inp, opt in zip(instruction, input_, output_):
|
||||
if inp is not None and inp != '':
|
||||
if inp.startswith('输入:'):
|
||||
inp = inp[3:]
|
||||
inst = f'{inst}\n{inp}'
|
||||
if opt is not None and opt != '':
|
||||
res.append(inst)
|
||||
output.append(opt)
|
||||
dataset = HfDataset.from_dict({'instruction': res, 'output': output})
|
||||
|
||||
def get_alpaca_gpt4_zh_dataset() -> HfDataset:
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'AI-ModelScope/alpaca-gpt4-data-zh', split='train').to_hf_dataset()
|
||||
|
||||
def _preprocess_input(inp: str) -> str:
|
||||
if inp.startswith('输入:'):
|
||||
inp = inp[3:]
|
||||
return inp
|
||||
|
||||
return _processing_alpaca(dataset, _preprocess_input)
|
||||
|
||||
|
||||
def get_finance_en_dataset() -> HfDataset:
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'wyj123456/finance_en', split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
|
||||
_multi_alpaca_language_list = [
|
||||
'ar', 'de', 'es', 'fr', 'id', 'ja', 'ko', 'pt', 'ru', 'th', 'vi'
|
||||
]
|
||||
|
||||
|
||||
def get_multi_alpaca(subset_name: str) -> HfDataset:
|
||||
"""
|
||||
subset_name:
|
||||
Language-key Language # examples
|
||||
ar Arabic 14,671
|
||||
de German 9,515
|
||||
es Spanish 9,958
|
||||
fr France 11,332
|
||||
id Indonesian 12,117
|
||||
ja Japanese 10,191
|
||||
ko Korean 14,402
|
||||
pt Portuguese 10,825
|
||||
ru Russian 14,286
|
||||
th Thai 11,496
|
||||
vi Vietnamese 13,908
|
||||
"""
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'damo/nlp_polylm_multialpaca_sft',
|
||||
subset_name=subset_name,
|
||||
split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
|
||||
def get_multi_alpaca_all() -> HfDataset:
|
||||
dataset_list = []
|
||||
for subset_name in _multi_alpaca_language_list:
|
||||
dataset = get_multi_alpaca(subset_name)
|
||||
dataset_list.append(dataset)
|
||||
dataset = concatenate_datasets(dataset_list)
|
||||
return dataset
|
||||
|
||||
|
||||
def get_alpaca_en_dataset() -> HfDataset:
|
||||
dataset_en: HfDataset = MsDataset.load(
|
||||
'AI-ModelScope/alpaca-gpt4-data-en', split='train').to_hf_dataset()
|
||||
dataset_en = dataset_en.remove_columns(['text'])
|
||||
return _processing_alpaca(dataset_en)
|
||||
def get_code_alpaca_en_dataset() -> HfDataset:
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'wyj123456/code_alpaca_en', split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
|
||||
def get_alpaca_zh_dataset() -> HfDataset:
|
||||
dataset_zh: HfDataset = MsDataset.load(
|
||||
'AI-ModelScope/alpaca-gpt4-data-zh', split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset_zh)
|
||||
def get_instinwild_zh_dataset():
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'wyj123456/instinwild', subset_name='default',
|
||||
split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
|
||||
def get_multi_alpaca_dataset() -> HfDataset:
|
||||
dataset_multi = []
|
||||
for subset_name in [
|
||||
'ar', 'de', 'es', 'fr', 'id', 'ja', 'ko', 'pt', 'ru', 'th', 'vi'
|
||||
]:
|
||||
dataset_sub: HfDataset = MsDataset.load(
|
||||
'damo/nlp_polylm_multialpaca_sft',
|
||||
subset_name=subset_name,
|
||||
split='train').to_hf_dataset()
|
||||
dataset_multi.append(dataset_sub)
|
||||
return _processing_multi_alpaca(dataset_multi)
|
||||
def get_instinwild_en_dataset():
|
||||
dataset: HfDataset = MsDataset.load(
|
||||
'wyj123456/instinwild', subset_name='subset',
|
||||
split='train').to_hf_dataset()
|
||||
return _processing_alpaca(dataset)
|
||||
|
||||
|
||||
def get_seed(random_state: RandomState) -> int:
|
||||
seed_max = np.iinfo(np.int32).max
|
||||
seed = random_state.randint(0, seed_max)
|
||||
return seed
|
||||
|
||||
|
||||
def process_dataset(dataset: HfDataset, dataset_test_size: float,
|
||||
dataset_sample: Optional[int],
|
||||
dataset_seed: int) -> Tuple[HfDataset, HfDataset]:
|
||||
random_state = np.random.RandomState(dataset_seed)
|
||||
if dataset_sample is not None:
|
||||
index = random_state.permutation(len(dataset))[:dataset_sample]
|
||||
dataset = dataset.select(index)
|
||||
dataset = dataset.train_test_split(
|
||||
dataset_test_size, seed=get_seed(random_state))
|
||||
return dataset['train'], dataset['test']
|
||||
|
||||
|
||||
DATASET_MAPPER = {
|
||||
'alpaca-en': get_alpaca_en_dataset,
|
||||
'alpaca-zh': get_alpaca_zh_dataset,
|
||||
'alpaca-multi': get_multi_alpaca_dataset,
|
||||
DATASET_MAPPING = {
|
||||
'alpaca-en': get_alpaca_gpt4_en_dataset,
|
||||
'alpaca-zh': get_alpaca_gpt4_zh_dataset,
|
||||
'finance-en': get_finance_en_dataset,
|
||||
'multi-alpaca-all': get_multi_alpaca_all,
|
||||
**{
|
||||
f'multi-alpaca-{k}': partial(get_multi_alpaca, k)
|
||||
for k in _multi_alpaca_language_list
|
||||
},
|
||||
'code-en': get_code_alpaca_en_dataset,
|
||||
'instinwild-zh': get_instinwild_zh_dataset,
|
||||
'instinwild-en': get_instinwild_en_dataset,
|
||||
}
|
||||
|
||||
|
||||
def get_dataset(dataset_name_list: List[str]) -> HfDataset:
|
||||
dataset_list = []
|
||||
for dataset_name in dataset_name_list:
|
||||
get_function = DATASET_MAPPER[dataset_name]
|
||||
get_function = DATASET_MAPPING[dataset_name]
|
||||
dataset_list.append(get_function())
|
||||
dataset = concatenate_datasets(dataset_list)
|
||||
return dataset
|
||||
|
||||
|
||||
def process_dataset(dataset: HfDataset, dataset_test_size: float,
|
||||
dataset_sample: int,
|
||||
dataset_seed: int) -> Tuple[HfDataset, HfDataset]:
|
||||
random_state = np.random.RandomState(dataset_seed)
|
||||
if dataset_sample >= 0:
|
||||
index = random_state.permutation(len(dataset))[:dataset_sample]
|
||||
dataset = dataset.select(index)
|
||||
dataset = dataset.train_test_split(
|
||||
dataset_test_size, seed=get_seed(random_state))
|
||||
return dataset['train'], dataset['test']
|
||||
|
||||
@@ -1,32 +1,37 @@
|
||||
import os
|
||||
from types import MethodType
|
||||
from typing import Any, Dict, NamedTuple, Optional
|
||||
|
||||
import torch
|
||||
from swift import get_logger
|
||||
from torch import dtype as Dtype
|
||||
|
||||
from modelscope import (AutoConfig, AutoModelForCausalLM, AutoTokenizer, Model,
|
||||
get_logger, read_config, snapshot_download)
|
||||
read_config, snapshot_download)
|
||||
from modelscope.models.nlp.chatglm2 import ChatGLM2Config, ChatGLM2Tokenizer
|
||||
from modelscope.models.nlp.qwen import QWenConfig, QWenTokenizer
|
||||
from modelscope.models.nlp.llama2 import Llama2Config, Llama2Tokenizer
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _add_special_token(tokenizer, special_token_mapper: Dict[str,
|
||||
Any]) -> None:
|
||||
for k, v in special_token_mapper:
|
||||
for k, v in special_token_mapper.items():
|
||||
setattr(tokenizer, k, v)
|
||||
assert tokenizer.eos_token is not None
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
|
||||
def get_model_tokenizer_default(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True):
|
||||
def get_model_tokenizer_from_repo(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True,
|
||||
model_config=None,
|
||||
**model_kwargs):
|
||||
"""load from an independent repository"""
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_dir, trust_remote_code=True)
|
||||
if model_config is None:
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_dir, trust_remote_code=True)
|
||||
model_config.torch_dtype = torch_dtype
|
||||
logger.info(f'model_config: {model_config}')
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
@@ -36,71 +41,98 @@ def get_model_tokenizer_default(model_dir: str,
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_dir,
|
||||
config=model_config,
|
||||
device_map='auto',
|
||||
torch_dtype=torch_dtype,
|
||||
trust_remote_code=True)
|
||||
trust_remote_code=True,
|
||||
**model_kwargs)
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def get_model_tokenizer_polylm(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True):
|
||||
"""load from an independent repository"""
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_dir, trust_remote_code=True)
|
||||
def get_model_tokenizer_from_sdk(config_class: type,
|
||||
tokenizer_class: type,
|
||||
model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True,
|
||||
model_config=None,
|
||||
**model_kwargs):
|
||||
"""load from ms library"""
|
||||
config = read_config(model_dir)
|
||||
logger.info(config)
|
||||
if model_config is None:
|
||||
model_config = config_class.from_pretrained(model_dir)
|
||||
model_config.torch_dtype = torch_dtype
|
||||
logger.info(f'model_config: {model_config}')
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False)
|
||||
logger.info(model_config)
|
||||
tokenizer = tokenizer_class.from_pretrained(model_dir)
|
||||
model = None
|
||||
if load_model:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model = Model.from_pretrained(
|
||||
model_dir,
|
||||
cfg_dict=config,
|
||||
config=model_config,
|
||||
device_map='auto',
|
||||
torch_dtype=torch_dtype,
|
||||
trust_remote_code=True)
|
||||
**model_kwargs)
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def get_model_tokenizer_baichuan13b(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True,
|
||||
**model_kwargs):
|
||||
# baichuan-13b does not implement the `get_input_embeddings` function
|
||||
model, tokenizer = get_model_tokenizer_from_repo(model_dir, torch_dtype,
|
||||
load_model,
|
||||
**model_kwargs)
|
||||
model.get_input_embeddings = MethodType(
|
||||
lambda self: self.model.embed_tokens, model)
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def get_model_tokenizer_chatglm2(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True):
|
||||
"""load from ms library"""
|
||||
config = read_config(model_dir)
|
||||
logger.info(config)
|
||||
model_config = ChatGLM2Config.from_pretrained(model_dir)
|
||||
model_config.torch_dtype = torch_dtype
|
||||
logger.info(model_config)
|
||||
tokenizer = ChatGLM2Tokenizer.from_pretrained(model_dir)
|
||||
model = None
|
||||
if load_model:
|
||||
model = Model.from_pretrained(
|
||||
model_dir,
|
||||
cfg_dict=config,
|
||||
config=model_config,
|
||||
device_map='auto',
|
||||
torch_dtype=torch_dtype)
|
||||
return model, tokenizer
|
||||
load_model: bool = True,
|
||||
**model_kwargs):
|
||||
if 'quantization_config' in model_kwargs:
|
||||
model_kwargs['quantization_config'].llm_int8_skip_modules = [
|
||||
'output_layer'
|
||||
]
|
||||
return get_model_tokenizer_from_sdk(ChatGLM2Config, ChatGLM2Tokenizer,
|
||||
model_dir, torch_dtype, load_model,
|
||||
**model_kwargs)
|
||||
|
||||
|
||||
def get_model_tokenizer_llama2(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True,
|
||||
**model_kwargs):
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_dir, trust_remote_code=True)
|
||||
model_config.pretraining_tp = 1
|
||||
return get_model_tokenizer_from_sdk(Llama2Config, Llama2Tokenizer,
|
||||
model_dir, torch_dtype, load_model,
|
||||
model_config, **model_kwargs)
|
||||
|
||||
|
||||
def get_model_tokenizer_qwen(model_dir: str,
|
||||
torch_dtype: Dtype,
|
||||
load_model: bool = True):
|
||||
config = read_config(model_dir)
|
||||
logger.info(config)
|
||||
model_config = QWenConfig.from_pretrained(model_dir)
|
||||
model_config.torch_dtype = torch_dtype
|
||||
logger.info(model_config)
|
||||
tokenizer = QWenTokenizer.from_pretrained(model_dir)
|
||||
model = None
|
||||
if load_model:
|
||||
model = Model.from_pretrained(
|
||||
model_dir,
|
||||
cfg_dict=config,
|
||||
config=model_config,
|
||||
device_map='auto',
|
||||
torch_dtype=torch_dtype)
|
||||
return model, tokenizer
|
||||
load_model: bool = True,
|
||||
**kwargs):
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_dir, trust_remote_code=True)
|
||||
mapper = {
|
||||
torch.float16: 'fp16',
|
||||
torch.bfloat16: 'bf16',
|
||||
torch.float32: 'fp32'
|
||||
}
|
||||
k_true = mapper[torch_dtype]
|
||||
for k in mapper.values():
|
||||
v = False
|
||||
if k == k_true:
|
||||
v = True
|
||||
setattr(model_config, k, v)
|
||||
|
||||
use_flash_attn = kwargs.pop('use_flash_attn', 'auto')
|
||||
model_config.use_flash_attn = use_flash_attn
|
||||
return get_model_tokenizer_from_repo(model_dir, torch_dtype, load_model,
|
||||
model_config, **kwargs)
|
||||
|
||||
|
||||
class LoRATM(NamedTuple):
|
||||
@@ -113,9 +145,9 @@ class LoRATM(NamedTuple):
|
||||
|
||||
|
||||
# Reference: 'https://modelscope.cn/models/{model_id}/summary'
|
||||
# keys: 'model_id', 'revision', 'torch_dtype', 'get_function',
|
||||
# keys: 'model_id', 'revision', 'get_function',
|
||||
# 'ignore_file_pattern', 'special_token_mapper', 'lora_TM'
|
||||
MODEL_MAPPER = {
|
||||
MODEL_MAPPING = {
|
||||
'baichuan-7b': {
|
||||
'model_id': 'baichuan-inc/baichuan-7B', # model id or model dir
|
||||
'revision': 'v1.0.7',
|
||||
@@ -124,69 +156,77 @@ MODEL_MAPPER = {
|
||||
'baichuan-13b': {
|
||||
'model_id': 'baichuan-inc/Baichuan-13B-Base',
|
||||
'revision': 'v1.0.3',
|
||||
'torch_dtype': torch.bfloat16,
|
||||
'get_function': get_model_tokenizer_baichuan13b,
|
||||
'lora_TM': LoRATM.baichuan
|
||||
},
|
||||
'chatglm2-6b': {
|
||||
'model_id': 'ZhipuAI/chatglm2-6b',
|
||||
'revision': 'v1.0.6',
|
||||
'revision': 'v1.0.7',
|
||||
'get_function': get_model_tokenizer_chatglm2,
|
||||
'lora_TM': LoRATM.chatglm2
|
||||
},
|
||||
'llama2-7b': {
|
||||
'model_id': 'modelscope/Llama-2-7b-ms',
|
||||
'revision': 'v1.0.2',
|
||||
'get_function': get_model_tokenizer_llama2,
|
||||
'ignore_file_pattern': [r'.+\.bin$'], # use safetensors
|
||||
'lora_TM': LoRATM.llama2
|
||||
},
|
||||
'llama2-13b': {
|
||||
'model_id': 'modelscope/Llama-2-13b-ms',
|
||||
'revision': 'v1.0.2',
|
||||
'get_function': get_model_tokenizer_llama2,
|
||||
'ignore_file_pattern': [r'.+\.bin$'],
|
||||
'lora_TM': LoRATM.llama2
|
||||
},
|
||||
'llama2-70b': {
|
||||
'model_id': 'modelscope/Llama-2-70b-ms',
|
||||
'revision': 'v1.0.0',
|
||||
'get_function': get_model_tokenizer_llama2,
|
||||
'ignore_file_pattern': [r'.+\.bin$'],
|
||||
'lora_TM': LoRATM.llama2
|
||||
},
|
||||
'openbuddy-llama2-13b': {
|
||||
'model_id': 'OpenBuddy/openbuddy-llama2-13b-v8.1-fp16',
|
||||
'revision': 'v1.0.0',
|
||||
'lora_TM': LoRATM.llama2
|
||||
'lora_TM': LoRATM.llama2,
|
||||
},
|
||||
'qwen-7b': {
|
||||
'model_id': 'QWen/qwen-7b',
|
||||
'revision': 'v1.0.0',
|
||||
'model_id': 'qwen/Qwen-7B',
|
||||
'revision': 'v.1.0.4',
|
||||
'get_function': get_model_tokenizer_qwen,
|
||||
'torch_dtype': torch.bfloat16,
|
||||
'lora_TM': LoRATM.qwen,
|
||||
},
|
||||
'polylm-13b': {
|
||||
'model_id': 'damo/nlp_polylm_13b_text_generation',
|
||||
'revision': 'v1.0.3',
|
||||
'get_function': get_model_tokenizer_polylm,
|
||||
'torch_dtype': torch.bfloat16,
|
||||
'lora_TM': LoRATM.polylm
|
||||
'special_token_mapper': {
|
||||
'eos_token': '<|endoftext|>'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def get_model_tokenizer(model_type: str,
|
||||
torch_dtype: Optional[Dtype] = None,
|
||||
load_model: bool = True):
|
||||
data = MODEL_MAPPER.get(model_type)
|
||||
load_model: bool = True,
|
||||
**kwargs):
|
||||
data = MODEL_MAPPING.get(model_type)
|
||||
if data is None:
|
||||
raise ValueError(f'model_type: {model_type}')
|
||||
|
||||
model_id = data['model_id']
|
||||
get_function = data.get('get_function', get_model_tokenizer_default)
|
||||
get_function = data.get('get_function', get_model_tokenizer_from_repo)
|
||||
ignore_file_pattern = data.get('ignore_file_pattern', [])
|
||||
special_token_mapper = data.get('special_token_mapper', {})
|
||||
if torch_dtype is None:
|
||||
torch_dtype = data.get('torch_dtype', torch.float16)
|
||||
|
||||
model_dir = model_id
|
||||
if not os.path.exists(model_id):
|
||||
revision = data.get('revision', 'master')
|
||||
model_dir = snapshot_download(
|
||||
model_id, revision, ignore_file_pattern=ignore_file_pattern)
|
||||
model_dir = kwargs.pop('model_dir', None)
|
||||
if model_dir is None:
|
||||
model_dir = model_id
|
||||
if not os.path.exists(model_id):
|
||||
revision = data.get('revision', 'master')
|
||||
model_dir = snapshot_download(
|
||||
model_id, revision, ignore_file_pattern=ignore_file_pattern)
|
||||
|
||||
model, tokenizer = get_function(model_dir, torch_dtype, load_model)
|
||||
model, tokenizer = get_function(model_dir, torch_dtype, load_model,
|
||||
**kwargs)
|
||||
_add_special_token(tokenizer, special_token_mapper)
|
||||
return model, tokenizer, model_dir
|
||||
|
||||
@@ -301,6 +301,7 @@ def inference(input_ids: List[int],
|
||||
print(f'{tag}{tokenizer.decode(input_ids)}', end='')
|
||||
input_ids = torch.tensor(input_ids)[None].cuda()
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
model.eval()
|
||||
generate_ids = model.generate(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
|
||||
@@ -5,47 +5,33 @@ import os
|
||||
import random
|
||||
import re
|
||||
import sys
|
||||
from functools import partial
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import json
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
#
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from matplotlib.axes import Axes
|
||||
from matplotlib.figure import Figure
|
||||
from numpy import ndarray
|
||||
from swift import LoRAConfig, Swift
|
||||
from tensorboard.backend.event_processing.event_accumulator import \
|
||||
EventAccumulator
|
||||
from torch import Tensor
|
||||
from torch import device as Device
|
||||
from torch import dtype as Dtype
|
||||
from torch.nn import Module
|
||||
from torch.nn.parameter import Parameter
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim import lr_scheduler as lrs
|
||||
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
||||
from torch.utils.data import Dataset
|
||||
#
|
||||
from torchmetrics import Accuracy, MeanMetric
|
||||
#
|
||||
from tqdm import tqdm
|
||||
|
||||
#
|
||||
from modelscope import (Model, MsDataset, get_logger, read_config,
|
||||
snapshot_download)
|
||||
from modelscope import Model, MsDataset, get_logger, read_config
|
||||
from modelscope.metrics.base import Metric
|
||||
from modelscope.metrics.builder import METRICS
|
||||
from modelscope.models.nlp.chatglm2 import ChatGLM2Tokenizer
|
||||
from modelscope.msdatasets.dataset_cls.custom_datasets import \
|
||||
TorchCustomDataset
|
||||
from modelscope.swift import LoRAConfig, Swift
|
||||
from modelscope.trainers import EpochBasedTrainer
|
||||
from modelscope.utils.config import Config, ConfigDict
|
||||
from modelscope.utils.config import ConfigDict
|
||||
from modelscope.utils.registry import default_group
|
||||
|
||||
#
|
||||
|
||||
@@ -209,8 +209,8 @@
|
||||
"LORA_ALPHA = 32\n",
|
||||
"LORA_DROPOUT_P = 0 # Arbitrary value\n",
|
||||
"lora_config = LoRAConfig(\n",
|
||||
" replace_modules=LORA_TARGET_MODULES,\n",
|
||||
" rank=LORA_RANK,\n",
|
||||
" target_modules=LORA_TARGET_MODULES,\n",
|
||||
" r=LORA_RANK,\n",
|
||||
" lora_alpha=LORA_ALPHA,\n",
|
||||
" lora_dropout=LORA_DROPOUT_P,\n",
|
||||
" pretrained_weights=CKPT_FAPTH)\n",
|
||||
|
||||
@@ -224,8 +224,8 @@
|
||||
"LORA_ALPHA = 32\n",
|
||||
"LORA_DROPOUT_P = 0.1\n",
|
||||
"lora_config = LoRAConfig(\n",
|
||||
" replace_modules=LORA_TARGET_MODULES,\n",
|
||||
" rank=LORA_RANK,\n",
|
||||
" target_modules=LORA_TARGET_MODULES,\n",
|
||||
" r=LORA_RANK,\n",
|
||||
" lora_alpha=LORA_ALPHA,\n",
|
||||
" lora_dropout=LORA_DROPOUT_P)\n",
|
||||
"logger.info(f'lora_config: {lora_config}')\n",
|
||||
|
||||
@@ -212,8 +212,8 @@
|
||||
"LORA_ALPHA = 32\n",
|
||||
"LORA_DROPOUT_P = 0 # Arbitrary value\n",
|
||||
"lora_config = LoRAConfig(\n",
|
||||
" replace_modules=LORA_TARGET_MODULES,\n",
|
||||
" rank=LORA_RANK,\n",
|
||||
" target_modules=LORA_TARGET_MODULES,\n",
|
||||
" r=LORA_RANK,\n",
|
||||
" lora_alpha=LORA_ALPHA,\n",
|
||||
" lora_dropout=LORA_DROPOUT_P,\n",
|
||||
" pretrained_weights=CKPT_FAPTH)\n",
|
||||
|
||||
@@ -234,8 +234,8 @@
|
||||
"LORA_ALPHA = 32\n",
|
||||
"LORA_DROPOUT_P = 0.1\n",
|
||||
"lora_config = LoRAConfig(\n",
|
||||
" replace_modules=LORA_TARGET_MODULES,\n",
|
||||
" rank=LORA_RANK,\n",
|
||||
" target_modules=LORA_TARGET_MODULES,\n",
|
||||
" r=LORA_RANK,\n",
|
||||
" lora_alpha=LORA_ALPHA,\n",
|
||||
" lora_dropout=LORA_DROPOUT_P)\n",
|
||||
"logger.info(f'lora_config: {lora_config}')\n",
|
||||
|
||||
@@ -2,8 +2,11 @@ import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
from modelscope import snapshot_download
|
||||
from modelscope.metainfo import Trainers
|
||||
from modelscope.models import Model
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.trainers import EpochBasedTrainer, build_trainer
|
||||
@@ -95,6 +98,12 @@ class StableDiffusionCustomArguments(TrainingArgs):
|
||||
'help': 'Path to json containing multiple concepts.',
|
||||
})
|
||||
|
||||
torch_type: str = field(
|
||||
default='float32',
|
||||
metadata={
|
||||
'help': ' The torch type, default is float32.',
|
||||
})
|
||||
|
||||
|
||||
training_args = StableDiffusionCustomArguments(
|
||||
task='text-to-image-synthesis').parse_cli()
|
||||
@@ -129,9 +138,18 @@ def cfg_modify_fn(cfg):
|
||||
return cfg
|
||||
|
||||
|
||||
# build model
|
||||
model_dir = snapshot_download(training_args.model)
|
||||
model = Model.from_pretrained(
|
||||
training_args.model,
|
||||
revision=args.model_revision,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32)
|
||||
|
||||
# build trainer and training
|
||||
kwargs = dict(
|
||||
model=training_args.model,
|
||||
model_revision=args.model_revision,
|
||||
model=model,
|
||||
cfg_file=os.path.join(model_dir, 'configuration.json'),
|
||||
class_prompt=args.class_prompt,
|
||||
instance_prompt=args.instance_prompt,
|
||||
modifier_token=args.modifier_token,
|
||||
@@ -148,9 +166,10 @@ kwargs = dict(
|
||||
work_dir=training_args.work_dir,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=validation_dataset,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32,
|
||||
cfg_modify_fn=cfg_modify_fn)
|
||||
|
||||
# build trainer and training
|
||||
trainer = build_trainer(name=Trainers.custom_diffusion, default_args=kwargs)
|
||||
trainer.train()
|
||||
|
||||
@@ -159,7 +178,7 @@ pipe = pipeline(
|
||||
task=Tasks.text_to_image_synthesis,
|
||||
model=training_args.model,
|
||||
custom_dir=training_args.work_dir + '/output',
|
||||
modifier_token='<new1>+<new2>',
|
||||
modifier_token=args.modifier_token,
|
||||
model_revision=args.model_revision)
|
||||
|
||||
output = pipe({'text': args.instance_prompt})
|
||||
|
||||
@@ -7,11 +7,12 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/custom/finetune_stable_d
|
||||
--class_data_dir './tmp/class_data' \
|
||||
--train_dataset_name 'buptwq/lora-stable-diffusion-finetune-dog' \
|
||||
--max_epochs 250 \
|
||||
--modifier_token "<new1>+<new2>" \
|
||||
--modifier_token "<new1>" \
|
||||
--num_class_images=200 \
|
||||
--save_ckpt_strategy 'by_epoch' \
|
||||
--logging_interval 1 \
|
||||
--train.dataloader.workers_per_gpu 0 \
|
||||
--evaluation.dataloader.workers_per_gpu 0 \
|
||||
--train.optimizer.lr 1e-5 \
|
||||
--torch_type 'float32' \
|
||||
--use_model_config true
|
||||
|
||||
@@ -2,11 +2,14 @@ import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
from modelscope import snapshot_download
|
||||
from modelscope.metainfo import Trainers
|
||||
from modelscope.models import Model
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.trainers import EpochBasedTrainer, build_trainer
|
||||
from modelscope.trainers import build_trainer
|
||||
from modelscope.trainers.training_args import TrainingArgs
|
||||
from modelscope.utils.constant import DownloadMode, Tasks
|
||||
|
||||
@@ -59,6 +62,12 @@ class StableDiffusionDreamboothArguments(TrainingArgs):
|
||||
'help': 'The pipeline prompt.',
|
||||
})
|
||||
|
||||
torch_type: str = field(
|
||||
default='float32',
|
||||
metadata={
|
||||
'help': ' The torch type, default is float32.',
|
||||
})
|
||||
|
||||
|
||||
training_args = StableDiffusionDreamboothArguments(
|
||||
task='text-to-image-synthesis').parse_cli()
|
||||
@@ -93,9 +102,18 @@ def cfg_modify_fn(cfg):
|
||||
return cfg
|
||||
|
||||
|
||||
# build model
|
||||
model_dir = snapshot_download(training_args.model)
|
||||
model = Model.from_pretrained(
|
||||
training_args.model,
|
||||
revision=args.model_revision,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32)
|
||||
|
||||
# build trainer and training
|
||||
kwargs = dict(
|
||||
model=training_args.model,
|
||||
model_revision=args.model_revision,
|
||||
model=model,
|
||||
cfg_file=os.path.join(model_dir, 'configuration.json'),
|
||||
work_dir=training_args.work_dir,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=validation_dataset,
|
||||
@@ -106,9 +124,10 @@ kwargs = dict(
|
||||
resolution=args.resolution,
|
||||
prior_loss_weight=args.prior_loss_weight,
|
||||
prompt=args.prompt,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32,
|
||||
cfg_modify_fn=cfg_modify_fn)
|
||||
|
||||
# build trainer and training
|
||||
trainer = build_trainer(
|
||||
name=Trainers.dreambooth_diffusion, default_args=kwargs)
|
||||
trainer.train()
|
||||
|
||||
@@ -17,4 +17,5 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/dreambooth/finetune_stab
|
||||
--train.dataloader.workers_per_gpu 0 \
|
||||
--evaluation.dataloader.workers_per_gpu 0 \
|
||||
--train.optimizer.lr 5e-6 \
|
||||
--torch_type 'float32' \
|
||||
--use_model_config true
|
||||
|
||||
@@ -2,11 +2,14 @@ import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
from modelscope import snapshot_download
|
||||
from modelscope.metainfo import Trainers
|
||||
from modelscope.models import Model
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.trainers import EpochBasedTrainer, build_trainer
|
||||
from modelscope.trainers import build_trainer
|
||||
from modelscope.trainers.training_args import TrainingArgs
|
||||
from modelscope.utils.constant import DownloadMode, Tasks
|
||||
|
||||
@@ -25,6 +28,12 @@ class StableDiffusionLoraArguments(TrainingArgs):
|
||||
'help': 'The rank size of lora intermediate linear.',
|
||||
})
|
||||
|
||||
torch_type: str = field(
|
||||
default='float32',
|
||||
metadata={
|
||||
'help': ' The torch type, default is float32.',
|
||||
})
|
||||
|
||||
|
||||
training_args = StableDiffusionLoraArguments(
|
||||
task='text-to-image-synthesis').parse_cli()
|
||||
@@ -59,16 +68,26 @@ def cfg_modify_fn(cfg):
|
||||
return cfg
|
||||
|
||||
|
||||
# build model
|
||||
model_dir = snapshot_download(training_args.model)
|
||||
model = Model.from_pretrained(
|
||||
training_args.model,
|
||||
revision=args.model_revision,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32)
|
||||
|
||||
# build trainer and training
|
||||
kwargs = dict(
|
||||
model=training_args.model,
|
||||
model_revision=args.model_revision,
|
||||
model=model,
|
||||
cfg_file=os.path.join(model_dir, 'configuration.json'),
|
||||
work_dir=training_args.work_dir,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=validation_dataset,
|
||||
lora_rank=args.lora_rank,
|
||||
torch_type=torch.float16
|
||||
if args.torch_type == 'float16' else torch.float32,
|
||||
cfg_modify_fn=cfg_modify_fn)
|
||||
|
||||
# build trainer and training
|
||||
trainer = build_trainer(name=Trainers.lora_diffusion, default_args=kwargs)
|
||||
trainer.train()
|
||||
|
||||
|
||||
@@ -5,10 +5,11 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/lora/finetune_stable_dif
|
||||
--work_dir './tmp/lora_diffusion' \
|
||||
--train_dataset_name 'buptwq/lora-stable-diffusion-finetune' \
|
||||
--max_epochs 100 \
|
||||
--lora_rank 4 \
|
||||
--lora_rank 16 \
|
||||
--save_ckpt_strategy 'by_epoch' \
|
||||
--logging_interval 1 \
|
||||
--train.dataloader.workers_per_gpu 0 \
|
||||
--evaluation.dataloader.workers_per_gpu 0 \
|
||||
--train.optimizer.lr 1e-4 \
|
||||
--torch_type 'float16' \
|
||||
--use_model_config true
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
|
||||
from modelscope.metainfo import Trainers
|
||||
from modelscope.msdatasets import MsDataset
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.trainers import build_trainer
|
||||
from modelscope.trainers.training_args import TrainingArgs
|
||||
from modelscope.utils.constant import DownloadMode, Tasks
|
||||
|
||||
|
||||
# Load configuration file and dataset
|
||||
@dataclass(init=False)
|
||||
class StableDiffusionXLLoraArguments(TrainingArgs):
|
||||
prompt: str = field(
|
||||
default='dog', metadata={
|
||||
'help': 'The pipeline prompt.',
|
||||
})
|
||||
|
||||
lora_rank: int = field(
|
||||
default=16,
|
||||
metadata={
|
||||
'help': 'The rank size of lora intermediate linear.',
|
||||
})
|
||||
|
||||
|
||||
training_args = StableDiffusionXLLoraArguments(
|
||||
task='text-to-image-synthesis').parse_cli()
|
||||
config, args = training_args.to_config()
|
||||
|
||||
if os.path.exists(args.train_dataset_name):
|
||||
# Load local dataset
|
||||
train_dataset = MsDataset.load(args.train_dataset_name)
|
||||
validation_dataset = MsDataset.load(args.train_dataset_name)
|
||||
else:
|
||||
# Load online dataset
|
||||
train_dataset = MsDataset.load(
|
||||
args.train_dataset_name,
|
||||
split='train',
|
||||
download_mode=DownloadMode.FORCE_REDOWNLOAD)
|
||||
validation_dataset = MsDataset.load(
|
||||
args.train_dataset_name,
|
||||
split='validation',
|
||||
download_mode=DownloadMode.FORCE_REDOWNLOAD)
|
||||
|
||||
|
||||
def cfg_modify_fn(cfg):
|
||||
if args.use_model_config:
|
||||
cfg.merge_from_dict(config)
|
||||
else:
|
||||
cfg = config
|
||||
cfg.train.lr_scheduler = {
|
||||
'type': 'LambdaLR',
|
||||
'lr_lambda': lambda _: 1,
|
||||
'last_epoch': -1
|
||||
}
|
||||
return cfg
|
||||
|
||||
|
||||
kwargs = dict(
|
||||
model=training_args.model,
|
||||
model_revision=args.model_revision,
|
||||
work_dir=training_args.work_dir,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=validation_dataset,
|
||||
lora_rank=args.lora_rank,
|
||||
cfg_modify_fn=cfg_modify_fn)
|
||||
|
||||
# build trainer and training
|
||||
trainer = build_trainer(name=Trainers.lora_diffusion_xl, default_args=kwargs)
|
||||
trainer.train()
|
||||
|
||||
# pipeline after training and save result
|
||||
pipe = pipeline(
|
||||
task=Tasks.text_to_image_synthesis,
|
||||
model=training_args.model,
|
||||
lora_dir=training_args.work_dir + '/output',
|
||||
model_revision=args.model_revision)
|
||||
|
||||
output = pipe({'text': args.prompt})
|
||||
# visualize the result on ipynb and save it
|
||||
output
|
||||
cv2.imwrite('./lora_xl_result.png', output['output_imgs'][0])
|
||||
@@ -0,0 +1,14 @@
|
||||
PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/lora_xl/finetune_stable_diffusion_xl_lora.py \
|
||||
--model 'AI-ModelScope/stable-diffusion-xl-base-1.0' \
|
||||
--model_revision 'v1.0.2' \
|
||||
--prompt "a dog" \
|
||||
--work_dir './tmp/lora_diffusion_xl' \
|
||||
--train_dataset_name 'buptwq/lora-stable-diffusion-finetune' \
|
||||
--max_epochs 100 \
|
||||
--lora_rank 16 \
|
||||
--save_ckpt_strategy 'by_epoch' \
|
||||
--logging_interval 1 \
|
||||
--train.dataloader.workers_per_gpu 0 \
|
||||
--evaluation.dataloader.workers_per_gpu 0 \
|
||||
--train.optimizer.lr 1e-4 \
|
||||
--use_model_config true
|
||||
@@ -26,8 +26,12 @@ if TYPE_CHECKING:
|
||||
from .pipelines import Pipeline, pipeline
|
||||
from .utils.hub import read_config, create_model_if_not_exist
|
||||
from .utils.logger import get_logger
|
||||
from .utils.constant import Tasks
|
||||
from .utils.hf_util import AutoConfig, GenerationConfig
|
||||
from .utils.hf_util import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM
|
||||
from .utils.hf_util import (AutoModel, AutoModelForCausalLM,
|
||||
AutoModelForSeq2SeqLM,
|
||||
AutoModelForSequenceClassification,
|
||||
AutoModelForTokenClassification)
|
||||
from .utils.hf_util import AutoTokenizer
|
||||
from .msdatasets import MsDataset
|
||||
|
||||
@@ -68,9 +72,12 @@ else:
|
||||
'pipelines': ['Pipeline', 'pipeline'],
|
||||
'utils.hub': ['read_config', 'create_model_if_not_exist'],
|
||||
'utils.logger': ['get_logger'],
|
||||
'utils.constant': ['Tasks'],
|
||||
'utils.hf_util': [
|
||||
'AutoConfig', 'GenerationConfig', 'AutoModel',
|
||||
'AutoModelForCausalLM', 'AutoModelForSeq2SeqLM', 'AutoTokenizer'
|
||||
'AutoModelForCausalLM', 'AutoModelForSeq2SeqLM', 'AutoTokenizer',
|
||||
'AutoModelForSequenceClassification',
|
||||
'AutoModelForTokenClassification'
|
||||
],
|
||||
'msdatasets': ['MsDataset']
|
||||
}
|
||||
|
||||
@@ -225,7 +225,11 @@ class Models(object):
|
||||
|
||||
clip_interrogator = 'clip-interrogator'
|
||||
stable_diffusion = 'stable-diffusion'
|
||||
stable_diffusion_xl = 'stable-diffusion-xl'
|
||||
videocomposer = 'videocomposer'
|
||||
text_to_360panorama_image = 'text-to-360panorama-image'
|
||||
image_to_video_model = 'image-to-video-model'
|
||||
video_to_video_model = 'video-to-video-model'
|
||||
|
||||
# science models
|
||||
unifold = 'unifold'
|
||||
@@ -241,6 +245,7 @@ class TaskModels(object):
|
||||
feature_extraction = 'feature-extraction'
|
||||
text_generation = 'text-generation'
|
||||
text_ranking = 'text-ranking'
|
||||
machine_reading_comprehension = 'machine-reading-comprehension'
|
||||
|
||||
|
||||
class Heads(object):
|
||||
@@ -496,6 +501,7 @@ class Pipelines(object):
|
||||
document_grounded_dialog_rerank = 'document-grounded-dialog-rerank'
|
||||
document_grounded_dialog_generate = 'document-grounded-dialog-generate'
|
||||
language_identification = 'language_identification'
|
||||
machine_reading_comprehension_for_ner = 'machine-reading-comprehension-for-ner'
|
||||
|
||||
# audio tasks
|
||||
sambert_hifigan_tts = 'sambert-hifigan-tts'
|
||||
@@ -536,6 +542,7 @@ class Pipelines(object):
|
||||
text_to_image_synthesis = 'text-to-image-synthesis'
|
||||
video_multi_modal_embedding = 'video-multi-modal-embedding'
|
||||
prost_text_video_retrieval = 'prost-text-video-retrieval'
|
||||
videocomposer = 'videocomposer'
|
||||
image_text_retrieval = 'image-text-retrieval'
|
||||
ofa_ocr_recognition = 'ofa-ocr-recognition'
|
||||
ofa_asr = 'ofa-asr'
|
||||
@@ -555,6 +562,8 @@ class Pipelines(object):
|
||||
efficient_diffusion_tuning = 'efficient-diffusion-tuning'
|
||||
multimodal_dialogue = 'multimodal-dialogue'
|
||||
llama2_text_generation_pipeline = 'llama2-text-generation-pipeline'
|
||||
image_to_video_task_pipeline = 'image-to-video-task-pipeline'
|
||||
video_to_video_pipeline = 'video-to-video-pipeline'
|
||||
|
||||
# science tasks
|
||||
protein_structure = 'unifold-protein-structure'
|
||||
@@ -952,6 +961,7 @@ class MultiModalTrainers(object):
|
||||
efficient_diffusion_tuning = 'efficient-diffusion-tuning'
|
||||
stable_diffusion = 'stable-diffusion'
|
||||
lora_diffusion = 'lora-diffusion'
|
||||
lora_diffusion_xl = 'lora-diffusion-xl'
|
||||
dreambooth_diffusion = 'dreambooth-diffusion'
|
||||
custom_diffusion = 'custom-diffusion'
|
||||
cones2_inference = 'cones2-inference'
|
||||
@@ -1083,6 +1093,7 @@ class Preprocessors(object):
|
||||
document_grounded_dialog_retrieval = 'document-grounded-dialog-retrieval'
|
||||
document_grounded_dialog_rerank = 'document-grounded-dialog-rerank'
|
||||
document_grounded_dialog_generate = 'document-grounded-dialog-generate'
|
||||
machine_reading_comprehension_for_ner = 'machine-reading-comprehension-for-ner'
|
||||
|
||||
# audio preprocessor
|
||||
linear_aec_fbank = 'linear-aec-fbank'
|
||||
|
||||
@@ -2,16 +2,10 @@
|
||||
|
||||
from typing import Dict
|
||||
|
||||
import numpy as np
|
||||
from sklearn.metrics import accuracy_score, f1_score
|
||||
|
||||
from modelscope.metainfo import Metrics
|
||||
from modelscope.outputs import OutputKeys
|
||||
from modelscope.utils.registry import default_group
|
||||
from modelscope.utils.tensor_utils import (torch_nested_detach,
|
||||
torch_nested_numpify)
|
||||
from .base import Metric
|
||||
from .builder import METRICS, MetricKeys
|
||||
from .builder import METRICS
|
||||
|
||||
|
||||
@METRICS.register_module(
|
||||
|
||||
@@ -4,10 +4,11 @@ import os.path as osp
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from modelscope.hub.check_model import check_local_model_is_latest
|
||||
from modelscope.hub.snapshot_download import snapshot_download
|
||||
from modelscope.metainfo import Tasks
|
||||
from modelscope.models.builder import build_backbone, build_model
|
||||
from modelscope.utils.automodel_utils import (can_load_by_ms,
|
||||
try_to_load_hf_model)
|
||||
from modelscope.utils.config import Config
|
||||
from modelscope.utils.constant import DEFAULT_MODEL_REVISION, Invoke, ModelFile
|
||||
from modelscope.utils.device import verify_device
|
||||
@@ -84,12 +85,22 @@ class Model(ABC):
|
||||
device(str, `optional`): The device to load the model.
|
||||
**kwargs:
|
||||
task(str, `optional`): The `Tasks` enumeration value to replace the task value
|
||||
read out of config in the `model_name_or_path`. This is useful when the model to be loaded is not
|
||||
equal to the model saved.
|
||||
For example, load a `backbone` into a `text-classification` model.
|
||||
Other kwargs will be directly fed into the `model` key, to replace the default configs.
|
||||
use_hf(bool): If set True, will use AutoModel in hf to initialize the model to keep compatibility
|
||||
with huggingface transformers.
|
||||
read out of config in the `model_name_or_path`. This is useful when the model to be loaded is not
|
||||
equal to the model saved.
|
||||
For example, load a `backbone` into a `text-classification` model.
|
||||
Other kwargs will be directly fed into the `model` key, to replace the default configs.
|
||||
use_hf(bool, `optional`):
|
||||
If set to True, it will initialize the model using AutoModel or AutoModelFor* from hf.
|
||||
If set to False, the model is loaded using the modelscope mode.
|
||||
If set to None, the loading mode will be automatically selected.
|
||||
ignore_file_pattern(List[str], `optional`):
|
||||
This parameter is passed to snapshot_download
|
||||
device_map(str | Dict[str, str], `optional`):
|
||||
This parameter is passed to AutoModel or AutoModelFor*
|
||||
torch_dtype(torch.dtype, `optional`):
|
||||
This parameter is passed to AutoModel or AutoModelFor*
|
||||
config(PretrainedConfig, `optional`):
|
||||
This parameter is passed to AutoModel or AutoModelFor*
|
||||
Returns:
|
||||
A model instance.
|
||||
|
||||
@@ -115,14 +126,14 @@ class Model(ABC):
|
||||
)
|
||||
|
||||
invoked_by = '%s/%s' % (Invoke.KEY, invoked_by)
|
||||
ignore_file_pattern = kwargs.get('ignore_file_pattern', None)
|
||||
local_model_dir = snapshot_download(
|
||||
model_name_or_path, revision, user_agent=invoked_by)
|
||||
model_name_or_path,
|
||||
revision,
|
||||
user_agent=invoked_by,
|
||||
ignore_file_pattern=ignore_file_pattern)
|
||||
logger.info(f'initialize model from {local_model_dir}')
|
||||
|
||||
if kwargs.pop('use_hf', False):
|
||||
from modelscope import AutoModel
|
||||
return AutoModel.from_pretrained(local_model_dir)
|
||||
|
||||
if cfg_dict is not None:
|
||||
cfg = cfg_dict
|
||||
else:
|
||||
@@ -134,6 +145,23 @@ class Model(ABC):
|
||||
model_cfg = cfg.model
|
||||
if hasattr(model_cfg, 'model_type') and not hasattr(model_cfg, 'type'):
|
||||
model_cfg.type = model_cfg.model_type
|
||||
model_type = model_cfg.type
|
||||
if isinstance(device, str) and device.startswith('gpu'):
|
||||
device = 'cuda' + device[3:]
|
||||
use_hf = kwargs.pop('use_hf', None)
|
||||
if use_hf is None and can_load_by_ms(local_model_dir, task_name,
|
||||
model_type):
|
||||
use_hf = False
|
||||
model = None
|
||||
if use_hf in {True, None}:
|
||||
model = try_to_load_hf_model(local_model_dir, task_name, use_hf,
|
||||
**kwargs)
|
||||
if model is not None:
|
||||
device_map = kwargs.get('device_map', None)
|
||||
if device_map is None and device is not None:
|
||||
model = model.to(device)
|
||||
return model
|
||||
# use ms
|
||||
model_cfg.model_dir = local_model_dir
|
||||
|
||||
# install and import remote repos before build
|
||||
|
||||
55
modelscope/models/cv/image_face_fusion/facegan/face_gan.py
Normal file
55
modelscope/models/cv/image_face_fusion/facegan/face_gan.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .gpen_model import FullGenerator
|
||||
|
||||
|
||||
class GPEN(object):
|
||||
|
||||
def __init__(self,
|
||||
model_path,
|
||||
size=512,
|
||||
channel_multiplier=2,
|
||||
device=torch.device('cpu')):
|
||||
self.mfile = model_path
|
||||
self.n_mlp = 8
|
||||
self.resolution = size
|
||||
self.device = device
|
||||
self.load_model(channel_multiplier)
|
||||
|
||||
def load_model(self, channel_multiplier=2):
|
||||
self.model = FullGenerator(self.resolution, 512, self.n_mlp,
|
||||
channel_multiplier).to(self.device)
|
||||
pretrained_dict = torch.load(self.mfile)
|
||||
self.model.load_state_dict(pretrained_dict)
|
||||
self.model.eval()
|
||||
|
||||
def process(self, im):
|
||||
preds = []
|
||||
imt = self.img2tensor(im)
|
||||
imt = F.interpolate(imt, (self.resolution, self.resolution))
|
||||
|
||||
with torch.no_grad():
|
||||
img_out, __ = self.model(imt)
|
||||
|
||||
face = self.tensor2img(img_out)
|
||||
|
||||
return face, preds
|
||||
|
||||
def img2tensor(self, img):
|
||||
img_t = torch.from_numpy(img).to(self.device)
|
||||
img_t = (img_t / 255. - 0.5) / 0.5
|
||||
img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB
|
||||
return img_t
|
||||
|
||||
def tensor2img(self, image_tensor, pmax=255.0, imtype=np.uint8):
|
||||
image_tensor = image_tensor * 0.5 + 0.5
|
||||
image_tensor = image_tensor.squeeze(0).permute(1, 2,
|
||||
0).flip(2) # RGB->BGR
|
||||
image_numpy = np.clip(image_tensor.float().cpu().numpy(), 0, 1) * pmax
|
||||
|
||||
return image_numpy.astype(imtype)
|
||||
@@ -1,93 +0,0 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
from .model import FullGenerator
|
||||
|
||||
|
||||
class GANWrap(object):
|
||||
|
||||
def __init__(self,
|
||||
model_path,
|
||||
size=256,
|
||||
channel_multiplier=1,
|
||||
device='cpu'):
|
||||
self.device = device
|
||||
self.mfile = model_path
|
||||
self.transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5),
|
||||
inplace=True),
|
||||
])
|
||||
self.batchSize = 2
|
||||
self.n_mlp = 8
|
||||
self.resolution = size
|
||||
self.load_model(channel_multiplier)
|
||||
|
||||
def load_model(self, channel_multiplier=2):
|
||||
self.model = FullGenerator(self.resolution, 512, self.n_mlp,
|
||||
channel_multiplier).to(self.device)
|
||||
pretrained_dict = torch.load(
|
||||
self.mfile, map_location=torch.device('cpu'))
|
||||
self.model.load_state_dict(pretrained_dict)
|
||||
self.model.eval()
|
||||
|
||||
def process_tensor(self, img_t, return_face=True):
|
||||
b, c, h, w = img_t.shape
|
||||
img_t = F.interpolate(img_t, (self.resolution, self.resolution))
|
||||
|
||||
with torch.no_grad():
|
||||
out, __ = self.model(img_t)
|
||||
|
||||
out = F.interpolate(out, (w, h))
|
||||
return out
|
||||
|
||||
def process(self, ims, return_face=True):
|
||||
res = []
|
||||
faces = []
|
||||
for i in range(0, len(ims), self.batchSize):
|
||||
sizes = []
|
||||
imt = None
|
||||
for im in ims[i:i + self.batchSize]:
|
||||
sizes.append(im.shape[0])
|
||||
im = cv2.resize(im, (self.resolution, self.resolution))
|
||||
im_pil = Image.fromarray(im)
|
||||
imt = self.img2tensor(im_pil) if imt is None else torch.cat(
|
||||
(imt, self.img2tensor(im_pil)), dim=0)
|
||||
|
||||
imt = torch.flip(imt, [1])
|
||||
with torch.no_grad():
|
||||
img_outs, __ = self.model(imt)
|
||||
|
||||
for sz, img_out in zip(sizes, img_outs):
|
||||
img = self.tensor2img(img_out)
|
||||
if return_face:
|
||||
faces.append(img)
|
||||
img = cv2.resize(img, (sz, sz), interpolation=cv2.INTER_AREA)
|
||||
res.append(img)
|
||||
|
||||
return res, faces
|
||||
|
||||
def img2tensor(self, img):
|
||||
img_t = self.transform(img).to(self.device)
|
||||
img_t = torch.unsqueeze(img_t, 0)
|
||||
return img_t
|
||||
|
||||
def tensor2img(self, image_tensor, bytes=255.0, imtype=np.uint8):
|
||||
if image_tensor.dim() == 3:
|
||||
image_numpy = image_tensor.cpu().float().numpy()
|
||||
else:
|
||||
image_numpy = image_tensor[0].cpu().float().numpy()
|
||||
image_numpy = np.transpose(image_numpy, (1, 2, 0))
|
||||
image_numpy = image_numpy[:, :, ::-1]
|
||||
image_numpy = np.clip(
|
||||
image_numpy * np.asarray([0.5, 0.5, 0.5])
|
||||
+ np.asarray([0.5, 0.5, 0.5]), 0, 1)
|
||||
image_numpy = image_numpy * bytes
|
||||
return image_numpy.astype(imtype)
|
||||
@@ -1,17 +1,19 @@
|
||||
# The implementation is adopted from stylegan2-pytorch,
|
||||
# made public available under the MIT License at https://github.com/rosinality/stylegan2-pytorch/blob/master/model.py
|
||||
# The implementation is adopted from InsightFace_Pytorch, made publicly available under the MIT License
|
||||
# at https://github.com/yangxy/GPEN
|
||||
import functools
|
||||
import itertools
|
||||
import math
|
||||
import operator
|
||||
import random
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.autograd import Function
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d
|
||||
|
||||
isconcat = True
|
||||
sss = 2 if isconcat else 1
|
||||
ratio = 2
|
||||
|
||||
|
||||
class PixelNorm(nn.Module):
|
||||
@@ -306,6 +308,7 @@ class NoiseInjection(nn.Module):
|
||||
def forward(self, image, noise=None):
|
||||
|
||||
if noise is not None:
|
||||
# print(image.shape, noise.shape)
|
||||
if isconcat:
|
||||
return torch.cat((image, self.weight * noise), dim=1) # concat
|
||||
return image + self.weight * noise
|
||||
@@ -356,7 +359,8 @@ class StyledConv(nn.Module):
|
||||
)
|
||||
|
||||
self.noise = NoiseInjection()
|
||||
self.activate = FusedLeakyReLU(out_channel * sss)
|
||||
feat_multiplier = 2
|
||||
self.activate = FusedLeakyReLU(out_channel * feat_multiplier)
|
||||
|
||||
def forward(self, input, style, noise=None):
|
||||
out = self.conv(input, style)
|
||||
@@ -405,12 +409,14 @@ class Generator(nn.Module):
|
||||
channel_multiplier=2,
|
||||
blur_kernel=[1, 3, 3, 1],
|
||||
lr_mlp=0.01,
|
||||
narrow=1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.size = size
|
||||
self.n_mlp = n_mlp
|
||||
self.style_dim = style_dim
|
||||
self.feat_multiplier = 2
|
||||
|
||||
layers = [PixelNorm()]
|
||||
|
||||
@@ -425,15 +431,16 @@ class Generator(nn.Module):
|
||||
self.style = nn.Sequential(*layers)
|
||||
|
||||
self.channels = {
|
||||
4: 512 // ratio,
|
||||
8: 512 // ratio,
|
||||
16: 512 // ratio,
|
||||
32: 512 // ratio,
|
||||
64: 256 // ratio * channel_multiplier,
|
||||
128: 128 // ratio * channel_multiplier,
|
||||
256: 64 // ratio * channel_multiplier,
|
||||
512: 32 // ratio * channel_multiplier,
|
||||
1024: 16 // ratio * channel_multiplier,
|
||||
4: int(512 * narrow),
|
||||
8: int(512 * narrow),
|
||||
16: int(512 * narrow),
|
||||
32: int(512 * narrow),
|
||||
64: int(256 * channel_multiplier * narrow),
|
||||
128: int(128 * channel_multiplier * narrow),
|
||||
256: int(64 * channel_multiplier * narrow),
|
||||
512: int(32 * channel_multiplier * narrow),
|
||||
1024: int(16 * channel_multiplier * narrow),
|
||||
2048: int(8 * channel_multiplier * narrow)
|
||||
}
|
||||
|
||||
self.input = ConstantInput(self.channels[4])
|
||||
@@ -443,7 +450,8 @@ class Generator(nn.Module):
|
||||
3,
|
||||
style_dim,
|
||||
blur_kernel=blur_kernel)
|
||||
self.to_rgb1 = ToRGB(self.channels[4] * sss, style_dim, upsample=False)
|
||||
self.to_rgb1 = ToRGB(
|
||||
self.channels[4] * self.feat_multiplier, style_dim, upsample=False)
|
||||
|
||||
self.log_size = int(math.log(size, 2))
|
||||
|
||||
@@ -458,23 +466,23 @@ class Generator(nn.Module):
|
||||
|
||||
self.convs.append(
|
||||
StyledConv(
|
||||
in_channel * sss,
|
||||
in_channel * self.feat_multiplier,
|
||||
out_channel,
|
||||
3,
|
||||
style_dim,
|
||||
upsample=True,
|
||||
blur_kernel=blur_kernel,
|
||||
))
|
||||
blur_kernel=blur_kernel))
|
||||
|
||||
self.convs.append(
|
||||
StyledConv(
|
||||
out_channel * sss,
|
||||
out_channel * self.feat_multiplier,
|
||||
out_channel,
|
||||
3,
|
||||
style_dim,
|
||||
blur_kernel=blur_kernel))
|
||||
|
||||
self.to_rgbs.append(ToRGB(out_channel * sss, style_dim))
|
||||
self.to_rgbs.append(
|
||||
ToRGB(out_channel * self.feat_multiplier, style_dim))
|
||||
|
||||
in_channel = out_channel
|
||||
|
||||
@@ -515,6 +523,9 @@ class Generator(nn.Module):
|
||||
styles = [self.style(s) for s in styles]
|
||||
|
||||
if noise is None:
|
||||
'''
|
||||
noise = [None] * (2 * (self.log_size - 2) + 1)
|
||||
'''
|
||||
noise = []
|
||||
batch = styles[0].shape[0]
|
||||
for i in range(self.n_mlp + 1):
|
||||
@@ -557,16 +568,14 @@ class Generator(nn.Module):
|
||||
skip = self.to_rgb1(out, latent[:, 1])
|
||||
|
||||
i = 1
|
||||
noise_i = 1
|
||||
|
||||
for conv1, conv2, to_rgb in zip(self.convs[::2], self.convs[1::2],
|
||||
self.to_rgbs):
|
||||
out = conv1(out, latent[:, i], noise=noise[(noise_i + 1) // 2])
|
||||
out = conv2(out, latent[:, i + 1], noise=noise[(noise_i + 2) // 2])
|
||||
for conv1, conv2, noise1, noise2, to_rgb in zip(
|
||||
self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2],
|
||||
self.to_rgbs):
|
||||
out = conv1(out, latent[:, i], noise=noise1)
|
||||
out = conv2(out, latent[:, i + 1], noise=noise2)
|
||||
skip = to_rgb(out, latent[:, i + 2], skip)
|
||||
|
||||
i += 2
|
||||
noise_i += 2
|
||||
|
||||
image = skip
|
||||
|
||||
@@ -652,21 +661,106 @@ class ResBlock(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
class FullGenerator(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size,
|
||||
style_dim,
|
||||
n_mlp,
|
||||
channel_multiplier=2,
|
||||
blur_kernel=[1, 3, 3, 1],
|
||||
lr_mlp=0.01,
|
||||
narrow=1,
|
||||
):
|
||||
super().__init__()
|
||||
channels = {
|
||||
4: int(512 * narrow),
|
||||
8: int(512 * narrow),
|
||||
16: int(512 * narrow),
|
||||
32: int(512 * narrow),
|
||||
64: int(256 * channel_multiplier * narrow),
|
||||
128: int(128 * channel_multiplier * narrow),
|
||||
256: int(64 * channel_multiplier * narrow),
|
||||
512: int(32 * channel_multiplier * narrow),
|
||||
1024: int(16 * channel_multiplier * narrow),
|
||||
2048: int(8 * channel_multiplier * narrow)
|
||||
}
|
||||
|
||||
self.log_size = int(math.log(size, 2))
|
||||
self.generator = Generator(
|
||||
size,
|
||||
style_dim,
|
||||
n_mlp,
|
||||
channel_multiplier=channel_multiplier,
|
||||
blur_kernel=blur_kernel,
|
||||
lr_mlp=lr_mlp,
|
||||
narrow=narrow)
|
||||
|
||||
conv = [ConvLayer(3, channels[size], 1)]
|
||||
self.ecd0 = nn.Sequential(*conv)
|
||||
in_channel = channels[size]
|
||||
|
||||
self.names = ['ecd%d' % i for i in range(self.log_size - 1)]
|
||||
for i in range(self.log_size, 2, -1):
|
||||
out_channel = channels[2**(i - 1)]
|
||||
conv = [ConvLayer(in_channel, out_channel, 3, downsample=True)]
|
||||
setattr(self, self.names[self.log_size - i + 1],
|
||||
nn.Sequential(*conv))
|
||||
in_channel = out_channel
|
||||
self.final_linear = nn.Sequential(
|
||||
EqualLinear(
|
||||
channels[4] * 4 * 4, style_dim, activation='fused_lrelu'))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
inputs,
|
||||
return_latents=False,
|
||||
inject_index=None,
|
||||
truncation=1,
|
||||
truncation_latent=None,
|
||||
input_is_latent=False,
|
||||
):
|
||||
noise = []
|
||||
for i in range(self.log_size - 1):
|
||||
ecd = getattr(self, self.names[i])
|
||||
inputs = ecd(inputs)
|
||||
noise.append(inputs)
|
||||
inputs = inputs.view(inputs.shape[0], -1)
|
||||
outs = self.final_linear(inputs)
|
||||
noise = list(
|
||||
itertools.chain.from_iterable(
|
||||
itertools.repeat(x, 2) for x in noise))[::-1]
|
||||
outs = self.generator([outs],
|
||||
return_latents,
|
||||
inject_index,
|
||||
truncation,
|
||||
truncation_latent,
|
||||
input_is_latent,
|
||||
noise=noise[1:])
|
||||
return outs
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
|
||||
def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
|
||||
def __init__(self,
|
||||
size,
|
||||
channel_multiplier=2,
|
||||
blur_kernel=[1, 3, 3, 1],
|
||||
narrow=1):
|
||||
super().__init__()
|
||||
|
||||
channels = {
|
||||
4: 512,
|
||||
8: 512,
|
||||
16: 512,
|
||||
32: 512,
|
||||
64: 256 * channel_multiplier,
|
||||
128: 128 * channel_multiplier,
|
||||
256: 64 * channel_multiplier,
|
||||
512: 32 * channel_multiplier,
|
||||
1024: 16 * channel_multiplier,
|
||||
4: int(512 * narrow),
|
||||
8: int(512 * narrow),
|
||||
16: int(512 * narrow),
|
||||
32: int(512 * narrow),
|
||||
64: int(256 * channel_multiplier * narrow),
|
||||
128: int(128 * channel_multiplier * narrow),
|
||||
256: int(64 * channel_multiplier * narrow),
|
||||
512: int(32 * channel_multiplier * narrow),
|
||||
1024: int(16 * channel_multiplier * narrow),
|
||||
2048: int(8 * channel_multiplier * narrow)
|
||||
}
|
||||
|
||||
convs = [ConvLayer(3, channels[size], 1)]
|
||||
@@ -713,48 +807,53 @@ class Discriminator(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
class FullGenerator(nn.Module):
|
||||
class FullGenerator_SR(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size,
|
||||
in_size,
|
||||
out_size,
|
||||
style_dim,
|
||||
n_mlp,
|
||||
channel_multiplier=2,
|
||||
blur_kernel=[1, 3, 3, 1],
|
||||
lr_mlp=0.01,
|
||||
narrow=1,
|
||||
):
|
||||
super().__init__()
|
||||
channels = {
|
||||
4: 512 // ratio,
|
||||
8: 512 // ratio,
|
||||
16: 512 // ratio,
|
||||
32: 512 // ratio,
|
||||
64: 256 // ratio * channel_multiplier,
|
||||
128: 128 // ratio * channel_multiplier,
|
||||
256: 64 // ratio * channel_multiplier,
|
||||
512: 32 // ratio * channel_multiplier,
|
||||
1024: 16 // ratio * channel_multiplier,
|
||||
4: int(512 * narrow),
|
||||
8: int(512 * narrow),
|
||||
16: int(512 * narrow),
|
||||
32: int(512 * narrow),
|
||||
64: int(256 * channel_multiplier * narrow),
|
||||
128: int(128 * channel_multiplier * narrow),
|
||||
256: int(64 * channel_multiplier * narrow),
|
||||
512: int(32 * channel_multiplier * narrow),
|
||||
1024: int(16 * channel_multiplier * narrow),
|
||||
2048: int(8 * channel_multiplier * narrow),
|
||||
}
|
||||
|
||||
self.log_size = int(math.log(size, 2))
|
||||
self.log_insize = int(math.log(in_size, 2))
|
||||
self.log_outsize = int(math.log(out_size, 2))
|
||||
self.generator = Generator(
|
||||
size,
|
||||
out_size,
|
||||
style_dim,
|
||||
n_mlp,
|
||||
channel_multiplier=channel_multiplier,
|
||||
blur_kernel=blur_kernel,
|
||||
lr_mlp=lr_mlp)
|
||||
lr_mlp=lr_mlp,
|
||||
narrow=narrow)
|
||||
|
||||
conv = [ConvLayer(3, channels[size], 1)]
|
||||
conv = [ConvLayer(3, channels[in_size], 1)]
|
||||
self.ecd0 = nn.Sequential(*conv)
|
||||
in_channel = channels[size]
|
||||
in_channel = channels[in_size]
|
||||
|
||||
self.names = ['ecd%d' % i for i in range(self.log_size - 1)]
|
||||
for i in range(self.log_size, 2, -1):
|
||||
self.names = ['ecd%d' % i for i in range(self.log_insize - 1)]
|
||||
for i in range(self.log_insize, 2, -1):
|
||||
out_channel = channels[2**(i - 1)]
|
||||
conv = [ConvLayer(in_channel, out_channel, 3, downsample=True)]
|
||||
setattr(self, self.names[self.log_size - i + 1],
|
||||
setattr(self, self.names[self.log_insize - i + 1],
|
||||
nn.Sequential(*conv))
|
||||
in_channel = out_channel
|
||||
self.final_linear = nn.Sequential(
|
||||
@@ -771,18 +870,22 @@ class FullGenerator(nn.Module):
|
||||
input_is_latent=False,
|
||||
):
|
||||
noise = []
|
||||
for i in range(self.log_size - 1):
|
||||
for i in range(self.log_outsize - self.log_insize):
|
||||
noise.append(None)
|
||||
for i in range(self.log_insize - 1):
|
||||
ecd = getattr(self, self.names[i])
|
||||
inputs = ecd(inputs)
|
||||
noise.append(inputs)
|
||||
inputs = inputs.view(inputs.shape[0], -1)
|
||||
outs = self.final_linear(inputs)
|
||||
outs = self.generator([outs],
|
||||
return_latents,
|
||||
inject_index,
|
||||
truncation,
|
||||
truncation_latent,
|
||||
input_is_latent,
|
||||
noise=noise[::-1])
|
||||
|
||||
return outs
|
||||
noise = list(
|
||||
itertools.chain.from_iterable(
|
||||
itertools.repeat(x, 2) for x in noise))[::-1]
|
||||
image, latent = self.generator([outs],
|
||||
return_latents,
|
||||
inject_index,
|
||||
truncation,
|
||||
truncation_latent,
|
||||
input_is_latent,
|
||||
noise=noise[1:])
|
||||
return image, latent
|
||||
@@ -1,4 +1,2 @@
|
||||
# The implementation is adopted from stylegan2-pytorch, made public available under the MIT License
|
||||
# at https://github.com/rosinality/stylegan2-pytorch
|
||||
from .fused_act import FusedLeakyReLU, fused_leaky_relu
|
||||
from .upfirdn2d import upfirdn2d
|
||||
|
||||
@@ -15,6 +15,77 @@ REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051],
|
||||
DEFAULT_CROP_SIZE = (96, 112)
|
||||
|
||||
|
||||
def _umeyama(src, dst, estimate_scale=True, scale=1.0):
|
||||
"""Estimate N-D similarity transformation with or without scaling.
|
||||
Parameters
|
||||
----------
|
||||
src : (M, N) array
|
||||
Source coordinates.
|
||||
dst : (M, N) array
|
||||
Destination coordinates.
|
||||
estimate_scale : bool
|
||||
Whether to estimate scaling factor.
|
||||
Returns
|
||||
-------
|
||||
T : (N + 1, N + 1)
|
||||
The homogeneous similarity transformation matrix. The matrix contains
|
||||
NaN values only if the problem is not well-conditioned.
|
||||
References
|
||||
----------
|
||||
.. [1] "Least-squares estimation of transformation parameters between two
|
||||
point patterns", Shinji Umeyama, PAMI 1991, :DOI:`10.1109/34.88573`
|
||||
"""
|
||||
|
||||
num = src.shape[0]
|
||||
dim = src.shape[1]
|
||||
|
||||
# Compute mean of src and dst.
|
||||
src_mean = src.mean(axis=0)
|
||||
dst_mean = dst.mean(axis=0)
|
||||
|
||||
# Subtract mean from src and dst.
|
||||
src_demean = src - src_mean
|
||||
dst_demean = dst - dst_mean
|
||||
|
||||
# Eq. (38).
|
||||
A = dst_demean.T @ src_demean / num
|
||||
|
||||
# Eq. (39).
|
||||
d = np.ones((dim, ), dtype=np.double)
|
||||
if np.linalg.det(A) < 0:
|
||||
d[dim - 1] = -1
|
||||
|
||||
T = np.eye(dim + 1, dtype=np.double)
|
||||
|
||||
U, S, V = np.linalg.svd(A)
|
||||
|
||||
# Eq. (40) and (43).
|
||||
rank = np.linalg.matrix_rank(A)
|
||||
if rank == 0:
|
||||
return np.nan * T
|
||||
elif rank == dim - 1:
|
||||
if np.linalg.det(U) * np.linalg.det(V) > 0:
|
||||
T[:dim, :dim] = U @ V
|
||||
else:
|
||||
s = d[dim - 1]
|
||||
d[dim - 1] = -1
|
||||
T[:dim, :dim] = U @ np.diag(d) @ V
|
||||
d[dim - 1] = s
|
||||
else:
|
||||
T[:dim, :dim] = U @ np.diag(d) @ V
|
||||
|
||||
if estimate_scale:
|
||||
# Eq. (41) and (42).
|
||||
scale = 1.0 / src_demean.var(axis=0).sum() * (S @ d)
|
||||
else:
|
||||
scale = scale
|
||||
|
||||
T[:dim, dim] = dst_mean - scale * (T[:dim, :dim] @ src_mean.T)
|
||||
T[:dim, :dim] *= scale
|
||||
|
||||
return T, scale
|
||||
|
||||
|
||||
class FaceWarpException(Exception):
|
||||
|
||||
def __str__(self):
|
||||
@@ -246,6 +317,65 @@ def warp_and_crop_face(src_img,
|
||||
return face_img
|
||||
|
||||
|
||||
def warp_and_crop_face_enhance(src_img,
|
||||
facial_pts,
|
||||
reference_pts=None,
|
||||
crop_size=(96, 112),
|
||||
align_type='smilarity'):
|
||||
if reference_pts is None:
|
||||
if crop_size[0] == 96 and crop_size[1] == 112:
|
||||
reference_pts = REFERENCE_FACIAL_POINTS
|
||||
else:
|
||||
default_square = False
|
||||
inner_padding_factor = 0
|
||||
outer_padding = (0, 0)
|
||||
output_size = crop_size
|
||||
|
||||
reference_pts = get_reference_facial_points(
|
||||
output_size, inner_padding_factor, outer_padding,
|
||||
default_square)
|
||||
|
||||
ref_pts = np.float32(reference_pts)
|
||||
ref_pts_shp = ref_pts.shape
|
||||
if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
|
||||
raise FaceWarpException(
|
||||
'reference_pts.shape must be (K,2) or (2,K) and K>2')
|
||||
|
||||
if ref_pts_shp[0] == 2:
|
||||
ref_pts = ref_pts.T
|
||||
|
||||
src_pts = np.float32(facial_pts)
|
||||
src_pts_shp = src_pts.shape
|
||||
if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
|
||||
raise FaceWarpException(
|
||||
'facial_pts.shape must be (K,2) or (2,K) and K>2')
|
||||
|
||||
if src_pts_shp[0] == 2:
|
||||
src_pts = src_pts.T
|
||||
|
||||
if src_pts.shape != ref_pts.shape:
|
||||
raise FaceWarpException(
|
||||
'facial_pts and reference_pts must have the same shape')
|
||||
|
||||
if align_type == 'cv2_affine':
|
||||
tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
|
||||
tfm_inv = cv2.getAffineTransform(ref_pts[0:3], src_pts[0:3])
|
||||
elif align_type == 'affine':
|
||||
tfm = get_affine_transform_matrix(src_pts, ref_pts)
|
||||
tfm_inv = get_affine_transform_matrix(ref_pts, src_pts)
|
||||
else:
|
||||
params, scale = _umeyama(src_pts, ref_pts)
|
||||
tfm = params[:2, :]
|
||||
|
||||
params, _ = _umeyama(ref_pts, src_pts, False, scale=1.0 / scale)
|
||||
tfm_inv = params[:2, :]
|
||||
|
||||
face_img = cv2.warpAffine(
|
||||
src_img, tfm, (crop_size[0], crop_size[1]), flags=3)
|
||||
|
||||
return face_img, tfm_inv
|
||||
|
||||
|
||||
def get_f5p(landmarks, np_img):
|
||||
eye_left = find_pupil(landmarks[36:41], np_img)
|
||||
eye_right = find_pupil(landmarks[42:47], np_img)
|
||||
|
||||
@@ -16,9 +16,10 @@ from modelscope.models.builder import MODELS
|
||||
from modelscope.models.cv.face_detection.peppa_pig_face.facer import FaceAna
|
||||
from modelscope.utils.constant import ModelFile, Tasks
|
||||
from modelscope.utils.logger import get_logger
|
||||
from .facegan.gan_wrap import GANWrap
|
||||
from .facegan.face_gan import GPEN
|
||||
from .facelib.align_trans import (get_f5p, get_reference_facial_points,
|
||||
warp_and_crop_face)
|
||||
warp_and_crop_face,
|
||||
warp_and_crop_face_enhance)
|
||||
from .network.aei_flow_net import AEI_Net
|
||||
from .network.bfm import ParametricFaceModel
|
||||
from .network.facerecon_model import ReconNetWrapper
|
||||
@@ -78,14 +79,6 @@ class ImageFaceFusion(TorchModel):
|
||||
self.face_model = ParametricFaceModel(bfm_folder=bfm_dir)
|
||||
self.face_model.to(self.device)
|
||||
|
||||
face_enhance_path = os.path.join(model_dir, 'faceEnhance',
|
||||
'350000-Ns256.pt')
|
||||
self.ganwrap = GANWrap(
|
||||
model_path=face_enhance_path,
|
||||
size=256,
|
||||
channel_multiplier=1,
|
||||
device=self.device)
|
||||
|
||||
self.facer = FaceAna(model_dir)
|
||||
|
||||
logger.info('load facefusion models done')
|
||||
@@ -94,6 +87,27 @@ class ImageFaceFusion(TorchModel):
|
||||
self.mask_init = cv2.resize(self.mask_init, (256, 256))
|
||||
self.mask = self.image_transform(self.mask_init, is_norm=False)
|
||||
|
||||
face_enhance_path = os.path.join(model_dir, 'faceEnhance',
|
||||
'GPEN-BFR-1024.pth')
|
||||
|
||||
if not os.path.exists(face_enhance_path):
|
||||
logger.warning(
|
||||
'model path not found, please update the latest model!')
|
||||
|
||||
self.ganwrap_1024 = GPEN(face_enhance_path, 1024, 2, self.device)
|
||||
|
||||
self.mask_enhance = np.zeros((512, 512), np.float32)
|
||||
cv2.rectangle(self.mask_enhance, (26, 26), (486, 486), (1, 1, 1), -1,
|
||||
cv2.LINE_AA)
|
||||
self.mask_enhance = cv2.GaussianBlur(self.mask_enhance, (101, 101), 11)
|
||||
self.mask_enhance = cv2.GaussianBlur(self.mask_enhance, (101, 101), 11)
|
||||
|
||||
default_square = True
|
||||
inner_padding_factor = 0.25
|
||||
outer_padding = (0, 0)
|
||||
self.reference_5pts_1024 = get_reference_facial_points(
|
||||
(1024, 1024), inner_padding_factor, outer_padding, default_square)
|
||||
|
||||
self.test_transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
@@ -157,7 +171,7 @@ class ImageFaceFusion(TorchModel):
|
||||
src_h, src_w, _ = img.shape
|
||||
boxes, landmarks, _ = self.facer.run(img)
|
||||
if boxes.shape[0] == 0:
|
||||
return None
|
||||
return None, None, None
|
||||
elif boxes.shape[0] > 1:
|
||||
max_area = 0
|
||||
max_index = 0
|
||||
@@ -168,9 +182,14 @@ class ImageFaceFusion(TorchModel):
|
||||
if area > max_area:
|
||||
max_index = i
|
||||
max_area = area
|
||||
return landmarks[max_index]
|
||||
|
||||
fw = boxes[max_index][2] - boxes[max_index][0]
|
||||
fh = boxes[max_index][3] - boxes[max_index][1]
|
||||
return landmarks[max_index], fw, fh
|
||||
else:
|
||||
return landmarks[0]
|
||||
fw = boxes[0][2] - boxes[0][0]
|
||||
fh = boxes[0][3] - boxes[0][1]
|
||||
return landmarks[0], fw, fh
|
||||
|
||||
def compute_3d_params(self, Xs, Xt):
|
||||
kp_fuse = {}
|
||||
@@ -198,6 +217,51 @@ class ImageFaceFusion(TorchModel):
|
||||
|
||||
return kp_fuse, kp_t
|
||||
|
||||
def process_enhance(self, im, f5p, fh, fw):
|
||||
height, width, _ = im.shape
|
||||
|
||||
of, tfm_inv = warp_and_crop_face_enhance(
|
||||
im,
|
||||
f5p,
|
||||
reference_pts=self.reference_5pts_1024,
|
||||
crop_size=(1024, 1024))
|
||||
ef, pred = self.ganwrap_1024.process(of)
|
||||
|
||||
tmp_mask = self.mask_enhance
|
||||
tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
|
||||
tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)
|
||||
|
||||
full_mask = np.zeros((height, width), dtype=np.float32)
|
||||
full_img = np.zeros(im.shape, dtype=np.uint8)
|
||||
|
||||
if min(fh, fw) < 40:
|
||||
ef = cv2.pyrDown(ef)
|
||||
ef = cv2.pyrDown(ef)
|
||||
ef = cv2.pyrUp(ef)
|
||||
ef = cv2.pyrUp(ef)
|
||||
elif min(fh, fw) < 60:
|
||||
ef = cv2.pyrDown(ef)
|
||||
ef = cv2.resize(ef, (0, 0), fx=2, fy=2)
|
||||
ef = cv2.resize(ef, (0, 0), fx=0.5, fy=0.5)
|
||||
ef = cv2.pyrUp(ef)
|
||||
elif min(fh, fw) < 80:
|
||||
ef = cv2.pyrDown(ef)
|
||||
ef = cv2.pyrUp(ef)
|
||||
elif min(fh, fw) < 100:
|
||||
ef = cv2.pyrDown(ef)
|
||||
ef = cv2.resize(ef, (0, 0), fx=2, fy=2)
|
||||
|
||||
tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
|
||||
|
||||
mask = tmp_mask - full_mask
|
||||
full_mask[np.where(mask > 0)] = tmp_mask[np.where(mask > 0)]
|
||||
full_img[np.where(mask > 0)] = tmp_img[np.where(mask > 0)]
|
||||
|
||||
full_mask = full_mask[:, :, np.newaxis]
|
||||
im = cv2.convertScaleAbs(im * (1 - full_mask) + full_img * full_mask)
|
||||
im = cv2.resize(im, (width, height))
|
||||
return im
|
||||
|
||||
def inference(self, template_img, user_img):
|
||||
ori_h, ori_w, _ = template_img.shape
|
||||
|
||||
@@ -205,14 +269,14 @@ class ImageFaceFusion(TorchModel):
|
||||
user_img = user_img.cpu().numpy()
|
||||
|
||||
user_img_bgr = user_img[:, :, ::-1]
|
||||
landmark_source = self.detect_face(user_img)
|
||||
landmark_source, _, _ = self.detect_face(user_img)
|
||||
if landmark_source is None:
|
||||
logger.warning('No face detected in user image!')
|
||||
return template_img
|
||||
f5p_user = get_f5p(landmark_source, user_img_bgr)
|
||||
|
||||
template_img_bgr = template_img[:, :, ::-1]
|
||||
landmark_template = self.detect_face(template_img)
|
||||
landmark_template, fw, fh = self.detect_face(template_img)
|
||||
if landmark_template is None:
|
||||
logger.warning('No face detected in template image!')
|
||||
return template_img
|
||||
@@ -235,7 +299,6 @@ class ImageFaceFusion(TorchModel):
|
||||
with torch.no_grad():
|
||||
kp_fuse, kp_t = self.compute_3d_params(Xs, Xt)
|
||||
Yt, _, _ = self.netG(Xt, Xs_embeds, kp_fuse, kp_t)
|
||||
Yt = self.ganwrap.process_tensor(Yt)
|
||||
Yt = Yt * 0.5 + 0.5
|
||||
Yt = torch.clamp(Yt, 0, 1)
|
||||
|
||||
@@ -247,6 +310,7 @@ class ImageFaceFusion(TorchModel):
|
||||
0).cpu().numpy()
|
||||
Yt_trans_inv = Yt_trans_inv.astype(np.float32)
|
||||
out_img = Yt_trans_inv[:, :, ::-1] * 255.
|
||||
out_img = self.process_enhance(out_img, f5p_template, fh, fw)
|
||||
|
||||
logger.info('model inference done')
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ if TYPE_CHECKING:
|
||||
from .team import TEAMForMultiModalSimilarity
|
||||
from .video_synthesis import TextToVideoSynthesis
|
||||
from .vldoc import VLDocForDocVLEmbedding
|
||||
from .videocomposer import VideoComposer
|
||||
|
||||
else:
|
||||
_import_structure = {
|
||||
@@ -44,6 +45,7 @@ else:
|
||||
'efficient_diffusion_tuning': ['EfficientStableDiffusion'],
|
||||
'mplug_owl': ['MplugOwlForConditionalGeneration'],
|
||||
'clip_interrogator': ['CLIP_Interrogator'],
|
||||
'videocomposer': ['VideoComposer'],
|
||||
}
|
||||
|
||||
import sys
|
||||
|
||||
@@ -13,21 +13,20 @@ from diffusers import (AutoencoderKL, DDPMScheduler, DiffusionPipeline,
|
||||
utils)
|
||||
from diffusers.models import cross_attention
|
||||
from diffusers.utils import deprecation_utils
|
||||
from swift import AdapterConfig, LoRAConfig, PromptConfig, Swift
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from modelscope import snapshot_download
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.models import TorchModel
|
||||
from modelscope.models.builder import MODELS
|
||||
from modelscope.models.multi_modal.efficient_diffusion_tuning.sd_lora import \
|
||||
LoRATuner
|
||||
from modelscope.outputs import OutputKeys
|
||||
from modelscope.swift import Swift
|
||||
from modelscope.swift.adapter import AdapterConfig
|
||||
from modelscope.swift.control_sd_lora import ControlLoRATuner
|
||||
from modelscope.swift.lora import LoRAConfig
|
||||
from modelscope.swift.prompt import PromptConfig
|
||||
from modelscope.swift.sd_lora import LoRATuner
|
||||
from modelscope.utils.checkpoint import save_checkpoint, save_configuration
|
||||
from modelscope.utils.config import Config
|
||||
from modelscope.utils.constant import ModelFile, Tasks
|
||||
from .control_sd_lora import ControlLoRATuner
|
||||
|
||||
utils.deprecate = lambda *arg, **kwargs: None
|
||||
deprecation_utils.deprecate = lambda *arg, **kwargs: None
|
||||
@@ -56,7 +55,10 @@ class EfficientStableDiffusion(TorchModel):
|
||||
super().__init__(model_dir, *args, **kwargs)
|
||||
tuner_name = kwargs.pop('tuner_name', 'lora')
|
||||
pretrained_model_name_or_path = kwargs.pop(
|
||||
'pretrained_model_name_or_path', 'runwayml/stable-diffusion-v1-5')
|
||||
'pretrained_model_name_or_path',
|
||||
'AI-ModelScope/stable-diffusion-v1-5')
|
||||
pretrained_model_name_or_path = snapshot_download(
|
||||
pretrained_model_name_or_path)
|
||||
tuner_config = kwargs.pop('tuner_config', None)
|
||||
pretrained_tuner = kwargs.pop('pretrained_tuner', None)
|
||||
revision = kwargs.pop('revision', None)
|
||||
|
||||
24
modelscope/models/multi_modal/image_to_video/__init__.py
Normal file
24
modelscope/models/multi_modal/image_to_video/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from modelscope.utils.import_utils import LazyImportModule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from .image_to_video_model import ImageToVideo
|
||||
|
||||
else:
|
||||
_import_structure = {
|
||||
'image_to_video_model': ['ImageToVideo'],
|
||||
}
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = LazyImportModule(
|
||||
__name__,
|
||||
globals()['__file__'],
|
||||
_import_structure,
|
||||
module_spec=__spec__,
|
||||
extra_objects={},
|
||||
)
|
||||
217
modelscope/models/multi_modal/image_to_video/image_to_video_model.py
Executable file
217
modelscope/models/multi_modal/image_to_video/image_to_video_model.py
Executable file
@@ -0,0 +1,217 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os
|
||||
import os.path as osp
|
||||
import random
|
||||
from copy import copy
|
||||
from typing import Any, Dict
|
||||
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
|
||||
import modelscope.models.multi_modal.image_to_video.utils.transforms as data
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.models.base import TorchModel
|
||||
from modelscope.models.builder import MODELS
|
||||
from modelscope.models.multi_modal.image_to_video.modules import *
|
||||
from modelscope.models.multi_modal.image_to_video.modules import (
|
||||
AutoencoderKL, FrozenOpenCLIPVisualEmbedder, Img2VidSDUNet)
|
||||
from modelscope.models.multi_modal.image_to_video.utils.config import cfg
|
||||
from modelscope.models.multi_modal.image_to_video.utils.diffusion import \
|
||||
GaussianDiffusion
|
||||
from modelscope.models.multi_modal.image_to_video.utils.seed import setup_seed
|
||||
from modelscope.models.multi_modal.image_to_video.utils.shedule import \
|
||||
beta_schedule
|
||||
from modelscope.utils.config import Config
|
||||
from modelscope.utils.constant import ModelFile, Tasks
|
||||
from modelscope.utils.device import create_device
|
||||
from modelscope.utils.logger import get_logger
|
||||
|
||||
__all__ = ['ImageToVideo']
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@MODELS.register_module(
|
||||
Tasks.image_to_video, module_name=Models.image_to_video_model)
|
||||
class ImageToVideo(TorchModel):
|
||||
r"""
|
||||
Image2Video aims to solve the task of generating high-definition videos based on input images.
|
||||
Image2Video is a video generation basic model developed by Alibaba Cloud, with a parameter size
|
||||
of approximately 2 billion. It has been pre trained on large-scale video and image data and
|
||||
fine-tuned on a small amount of high-quality data. The data is widely distributed and diverse
|
||||
in categories, and the model has good generalization ability for different types of data
|
||||
|
||||
Paper link: https://arxiv.org/abs/2306.02018
|
||||
|
||||
Attributes:
|
||||
diffusion: diffusion model for DDIM.
|
||||
autoencoder: decode the latent representation into visual space.
|
||||
clip_encoder: encode the image into image embedding.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir, *args, **kwargs):
|
||||
r"""
|
||||
Args:
|
||||
model_dir (`str` or `os.PathLike`)
|
||||
Can be either:
|
||||
- A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co
|
||||
or modelscope.cn. Valid model ids can be located at the root-level, like `bert-base-uncased`,
|
||||
or namespaced under a user or organization name, like `dbmdz/bert-base-german-cased`.
|
||||
- A path to a *directory* containing model weights saved using
|
||||
[`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
|
||||
- A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
|
||||
this case, `from_tf` should be set to `True` and a configuration object should be provided as
|
||||
`config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
|
||||
PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
|
||||
- A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g,
|
||||
`./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to
|
||||
`True`.
|
||||
"""
|
||||
super().__init__(model_dir=model_dir, *args, **kwargs)
|
||||
|
||||
self.config = Config.from_file(
|
||||
osp.join(model_dir, ModelFile.CONFIGURATION))
|
||||
|
||||
# assign default value
|
||||
cfg.batch_size = self.config.model.model_cfg.batch_size
|
||||
cfg.target_fps = self.config.model.model_cfg.target_fps
|
||||
cfg.max_frames = self.config.model.model_cfg.max_frames
|
||||
cfg.latent_hei = self.config.model.model_cfg.latent_hei
|
||||
cfg.latent_wid = self.config.model.model_cfg.latent_wid
|
||||
cfg.model_path = osp.join(model_dir,
|
||||
self.config.model.model_args.ckpt_unet)
|
||||
|
||||
required_device = kwargs.pop('device', 'gpu')
|
||||
self.device = create_device(required_device)
|
||||
|
||||
if 'seed' in self.config.model.model_args.keys():
|
||||
cfg.seed = self.config.model.model_args.seed
|
||||
else:
|
||||
cfg.seed = random.randint(0, 99999)
|
||||
setup_seed(cfg.seed)
|
||||
|
||||
# transform
|
||||
vid_trans = data.Compose([
|
||||
data.CenterCropWide(size=(cfg.resolution[0], cfg.resolution[0])),
|
||||
data.Resize(cfg.vit_resolution),
|
||||
data.ToTensor(),
|
||||
data.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)
|
||||
])
|
||||
self.vid_trans = vid_trans
|
||||
|
||||
cfg.embedder.pretrained = osp.join(
|
||||
model_dir, self.config.model.model_args.ckpt_clip)
|
||||
clip_encoder = FrozenOpenCLIPVisualEmbedder(
|
||||
device=self.device, **cfg.embedder)
|
||||
clip_encoder.model.to(self.device)
|
||||
self.clip_encoder = clip_encoder
|
||||
logger.info(f'Build encoder with {cfg.embedder.type}')
|
||||
|
||||
# [unet]
|
||||
generator = Img2VidSDUNet(**cfg.UNet)
|
||||
generator = generator.to(self.device)
|
||||
generator.eval()
|
||||
load_dict = torch.load(cfg.model_path, map_location='cpu')
|
||||
ret = generator.load_state_dict(load_dict['state_dict'], strict=True)
|
||||
self.generator = generator
|
||||
logger.info('Load model {} path {}, with local status {}'.format(
|
||||
cfg.UNet.type, cfg.model_path, ret))
|
||||
|
||||
# [diffusion]
|
||||
betas = beta_schedule(
|
||||
'linear_sd',
|
||||
cfg.num_timesteps,
|
||||
init_beta=0.00085,
|
||||
last_beta=0.0120)
|
||||
diffusion = GaussianDiffusion(
|
||||
betas=betas,
|
||||
mean_type=cfg.mean_type,
|
||||
var_type=cfg.var_type,
|
||||
loss_type=cfg.loss_type,
|
||||
rescale_timesteps=False,
|
||||
noise_strength=getattr(cfg, 'noise_strength', 0))
|
||||
self.diffusion = diffusion
|
||||
logger.info('Build diffusion with type of GaussianDiffusion')
|
||||
|
||||
# [auotoencoder]
|
||||
cfg.auto_encoder.pretrained = osp.join(
|
||||
model_dir, self.config.model.model_args.ckpt_autoencoder)
|
||||
autoencoder = AutoencoderKL(**cfg.auto_encoder)
|
||||
autoencoder.eval()
|
||||
for param in autoencoder.parameters():
|
||||
param.requires_grad = False
|
||||
autoencoder.to(self.device)
|
||||
self.autoencoder = autoencoder
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
zero_feature = torch.zeros(1, 1, cfg.UNet.input_dim).to(self.device)
|
||||
self.zero_feature = zero_feature
|
||||
self.fps_tensor = torch.tensor([cfg.target_fps],
|
||||
dtype=torch.long,
|
||||
device=self.device)
|
||||
self.cfg = cfg
|
||||
|
||||
def forward(self, input: Dict[str, Any]):
|
||||
r"""
|
||||
The entry function of image to video task.
|
||||
1. Using diffusion model to generate the video's latent representation.
|
||||
2. Using autoencoder to decode the video's latent representation to visual space.
|
||||
|
||||
Args:
|
||||
input (`Dict[Str, Any]`):
|
||||
The input of the task
|
||||
Returns:
|
||||
A generated video (as pytorch tensor).
|
||||
"""
|
||||
|
||||
vit_frame = input['vit_frame']
|
||||
cfg = self.cfg
|
||||
|
||||
img_embedding = self.clip_encoder(vit_frame).unsqueeze(1)
|
||||
|
||||
noise = self.build_noise()
|
||||
zero_feature = copy(self.zero_feature)
|
||||
with torch.no_grad():
|
||||
with amp.autocast(enabled=cfg.use_fp16):
|
||||
model_kwargs = [{
|
||||
'y': img_embedding,
|
||||
'fps': self.fps_tensor
|
||||
}, {
|
||||
'y': zero_feature.repeat(cfg.batch_size, 1, 1),
|
||||
'fps': self.fps_tensor
|
||||
}]
|
||||
gen_video = self.diffusion.ddim_sample_loop(
|
||||
noise=noise,
|
||||
model=self.generator,
|
||||
model_kwargs=model_kwargs,
|
||||
guide_scale=cfg.guide_scale,
|
||||
ddim_timesteps=cfg.ddim_timesteps,
|
||||
eta=0.0)
|
||||
|
||||
gen_video = 1. / cfg.scale_factor * gen_video
|
||||
gen_video = rearrange(gen_video, 'b c f h w -> (b f) c h w')
|
||||
chunk_size = min(cfg.decoder_bs, gen_video.shape[0])
|
||||
gen_video_list = torch.chunk(
|
||||
gen_video, gen_video.shape[0] // chunk_size, dim=0)
|
||||
decode_generator = []
|
||||
for vd_data in gen_video_list:
|
||||
gen_frames = self.autoencoder.decode(vd_data)
|
||||
decode_generator.append(gen_frames)
|
||||
|
||||
gen_video = torch.cat(decode_generator, dim=0)
|
||||
gen_video = rearrange(
|
||||
gen_video, '(b f) c h w -> b c f h w', b=cfg.batch_size)
|
||||
|
||||
return gen_video.type(torch.float32).cpu()
|
||||
|
||||
def build_noise(self):
|
||||
cfg = self.cfg
|
||||
noise = torch.randn(
|
||||
[1, 4, cfg.max_frames, cfg.latent_hei,
|
||||
cfg.latent_wid]).to(self.device)
|
||||
if cfg.noise_strength > 0:
|
||||
b, c, f, *_ = noise.shape
|
||||
offset_noise = torch.randn(b, c, f, 1, 1, device=noise.device)
|
||||
noise = noise + cfg.noise_strength * offset_noise
|
||||
return noise.contiguous()
|
||||
5
modelscope/models/multi_modal/image_to_video/modules/__init__.py
Executable file
5
modelscope/models/multi_modal/image_to_video/modules/__init__.py
Executable file
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .autoencoder import *
|
||||
from .embedder import *
|
||||
from .unet_i2v import *
|
||||
576
modelscope/models/multi_modal/image_to_video/modules/autoencoder.py
Executable file
576
modelscope/models/multi_modal/image_to_video/modules/autoencoder.py
Executable file
@@ -0,0 +1,576 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import collections
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modelscope.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def nonlinearity(x):
|
||||
# swish
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
def Normalize(in_channels, num_groups=32):
|
||||
return torch.nn.GroupNorm(
|
||||
num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
|
||||
class DiagonalGaussianDistribution(object):
|
||||
|
||||
def __init__(self, parameters, deterministic=False):
|
||||
self.parameters = parameters
|
||||
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
|
||||
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
|
||||
self.deterministic = deterministic
|
||||
self.std = torch.exp(0.5 * self.logvar)
|
||||
self.var = torch.exp(self.logvar)
|
||||
if self.deterministic:
|
||||
self.var = self.std = torch.zeros_like(
|
||||
self.mean).to(device=self.parameters.device)
|
||||
|
||||
def sample(self):
|
||||
x = self.mean + self.std * torch.randn(
|
||||
self.mean.shape).to(device=self.parameters.device)
|
||||
return x
|
||||
|
||||
def kl(self, other=None):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
else:
|
||||
if other is None:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
|
||||
dim=[1, 2, 3])
|
||||
else:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean - other.mean, 2) / other.var
|
||||
+ self.var / other.var - 1.0 - self.logvar + other.logvar,
|
||||
dim=[1, 2, 3])
|
||||
|
||||
def nll(self, sample, dims=[1, 2, 3]):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
logtwopi = np.log(2.0 * np.pi)
|
||||
return 0.5 * torch.sum(
|
||||
logtwopi + self.logvar
|
||||
+ torch.pow(sample - self.mean, 2) / self.var,
|
||||
dim=dims)
|
||||
|
||||
def mode(self):
|
||||
return self.mean
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
in_channels,
|
||||
out_channels=None,
|
||||
conv_shortcut=False,
|
||||
dropout,
|
||||
temb_channels=512):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_conv_shortcut = conv_shortcut
|
||||
|
||||
self.norm1 = Normalize(in_channels)
|
||||
self.conv1 = torch.nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if temb_channels > 0:
|
||||
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
|
||||
self.norm2 = Normalize(out_channels)
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = torch.nn.Conv2d(
|
||||
out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
self.conv_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
else:
|
||||
self.nin_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0)
|
||||
|
||||
def forward(self, x, temb):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
if temb is not None:
|
||||
h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
|
||||
|
||||
h = self.norm2(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.dropout(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
x = self.conv_shortcut(x)
|
||||
else:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = Normalize(in_channels)
|
||||
self.q = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.k = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.v = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.proj_out = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h * w)
|
||||
q = q.permute(0, 2, 1)
|
||||
k = k.reshape(b, c, h * w)
|
||||
w_ = torch.bmm(q, k)
|
||||
w_ = w_ * (int(c)**(-0.5))
|
||||
w_ = torch.nn.functional.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h * w)
|
||||
w_ = w_.permute(0, 2, 1)
|
||||
h_ = torch.bmm(v, w_)
|
||||
h_ = h_.reshape(b, c, h, w)
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x + h_
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.nn.functional.interpolate(
|
||||
x, scale_factor=2.0, mode='nearest')
|
||||
if self.with_conv:
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
if self.with_conv:
|
||||
pad = (0, 1, 0, 1)
|
||||
x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
|
||||
x = self.conv(x)
|
||||
else:
|
||||
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
double_z=True,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignore_kwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
|
||||
# downsampling
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1, ) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in,
|
||||
2 * z_channels if double_z else z_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1], temb)
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
give_pre_end=False,
|
||||
tanh_out=False,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignorekwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.give_pre_end = give_pre_end
|
||||
self.tanh_out = tanh_out
|
||||
|
||||
# compute block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2**(self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
logger.info('Working with z of shape {} = {} dimensions.'.format(
|
||||
self.z_shape, np.prod(self.z_shape)))
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up)
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z):
|
||||
self.last_z_shape = z.shape
|
||||
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h, temb)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
if self.give_pre_end:
|
||||
return h
|
||||
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
if self.tanh_out:
|
||||
h = torch.tanh(h)
|
||||
return h
|
||||
|
||||
|
||||
class AutoencoderKL(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
embed_dim,
|
||||
pretrained=None,
|
||||
ignore_keys=[],
|
||||
image_key='image',
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
ema_decay=None,
|
||||
learn_logvar=False,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
self.learn_logvar = learn_logvar
|
||||
self.image_key = image_key
|
||||
self.encoder = Encoder(**ddconfig)
|
||||
self.decoder = Decoder(**ddconfig)
|
||||
assert ddconfig['double_z']
|
||||
self.quant_conv = torch.nn.Conv2d(2 * ddconfig['z_channels'],
|
||||
2 * embed_dim, 1)
|
||||
self.post_quant_conv = torch.nn.Conv2d(embed_dim,
|
||||
ddconfig['z_channels'], 1)
|
||||
self.embed_dim = embed_dim
|
||||
if colorize_nlabels is not None:
|
||||
assert type(colorize_nlabels) == int
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, colorize_nlabels, 1, 1))
|
||||
if monitor is not None:
|
||||
self.monitor = monitor
|
||||
|
||||
self.use_ema = ema_decay is not None
|
||||
|
||||
if pretrained is not None:
|
||||
self.init_from_ckpt(pretrained, ignore_keys=ignore_keys)
|
||||
|
||||
def init_from_ckpt(self, path, ignore_keys=list()):
|
||||
sd = torch.load(path, map_location='cpu')['state_dict']
|
||||
keys = list(sd.keys())
|
||||
sd_new = collections.OrderedDict()
|
||||
for k in keys:
|
||||
if k.find('first_stage_model') >= 0:
|
||||
k_new = k.split('first_stage_model.')[-1]
|
||||
sd_new[k_new] = sd[k]
|
||||
self.load_state_dict(sd_new, strict=True)
|
||||
logger.info(f'Restored from {path}')
|
||||
|
||||
def on_train_batch_end(self, *args, **kwargs):
|
||||
if self.use_ema:
|
||||
self.model_ema(self)
|
||||
|
||||
def encode(self, x):
|
||||
h = self.encoder(x)
|
||||
moments = self.quant_conv(h)
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
return posterior
|
||||
|
||||
def decode(self, z):
|
||||
z = self.post_quant_conv(z)
|
||||
dec = self.decoder(z)
|
||||
return dec
|
||||
|
||||
def forward(self, input, sample_posterior=True):
|
||||
posterior = self.encode(input)
|
||||
if sample_posterior:
|
||||
z = posterior.sample()
|
||||
else:
|
||||
z = posterior.mode()
|
||||
dec = self.decode(z)
|
||||
return dec, posterior
|
||||
|
||||
def get_input(self, batch, k):
|
||||
x = batch[k]
|
||||
if len(x.shape) == 3:
|
||||
x = x[..., None]
|
||||
x = x.permute(0, 3, 1,
|
||||
2).to(memory_format=torch.contiguous_format).float()
|
||||
return x
|
||||
|
||||
def get_last_layer(self):
|
||||
return self.decoder.conv_out.weight
|
||||
|
||||
@torch.no_grad()
|
||||
def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
if not only_inputs:
|
||||
xrec, posterior = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec.shape[1] > 3
|
||||
x = self.to_rgb(x)
|
||||
xrec = self.to_rgb(xrec)
|
||||
log['samples'] = self.decode(torch.randn_like(posterior.sample()))
|
||||
log['reconstructions'] = xrec
|
||||
if log_ema or self.use_ema:
|
||||
with self.ema_scope():
|
||||
xrec_ema, posterior_ema = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec_ema.shape[1] > 3
|
||||
xrec_ema = self.to_rgb(xrec_ema)
|
||||
log['samples_ema'] = self.decode(
|
||||
torch.randn_like(posterior_ema.sample()))
|
||||
log['reconstructions_ema'] = xrec_ema
|
||||
log['inputs'] = x
|
||||
return log
|
||||
|
||||
def to_rgb(self, x):
|
||||
assert self.image_key == 'segmentation'
|
||||
if not hasattr(self, 'colorize'):
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, x.shape[1], 1, 1).to(x))
|
||||
x = F.conv2d(x, weight=self.colorize)
|
||||
x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
|
||||
return x
|
||||
|
||||
|
||||
class IdentityFirstStage(torch.nn.Module):
|
||||
|
||||
def __init__(self, *args, vq_interface=False, **kwargs):
|
||||
self.vq_interface = vq_interface
|
||||
super().__init__()
|
||||
|
||||
def encode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def decode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def quantize(self, x, *args, **kwargs):
|
||||
if self.vq_interface:
|
||||
return x, None, [None, None, None]
|
||||
return x
|
||||
|
||||
def forward(self, x, *args, **kwargs):
|
||||
return x
|
||||
81
modelscope/models/multi_modal/image_to_video/modules/embedder.py
Executable file
81
modelscope/models/multi_modal/image_to_video/modules/embedder.py
Executable file
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as T
|
||||
|
||||
|
||||
class FrozenOpenCLIPVisualEmbedder(nn.Module):
|
||||
"""
|
||||
Uses the OpenCLIP transformer encoder for text
|
||||
"""
|
||||
LAYERS = ['last', 'penultimate']
|
||||
|
||||
def __init__(self,
|
||||
pretrained,
|
||||
vit_resolution=(224, 224),
|
||||
arch='ViT-H-14',
|
||||
device='cuda',
|
||||
max_length=77,
|
||||
freeze=True,
|
||||
layer='last',
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
assert layer in self.LAYERS
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=pretrained)
|
||||
|
||||
del model.transformer
|
||||
self.model = model
|
||||
data_white = np.ones(
|
||||
(vit_resolution[0], vit_resolution[1], 3), dtype=np.uint8) * 255
|
||||
self.white_image = preprocess(T.ToPILImage()(data_white)).unsqueeze(0)
|
||||
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
if freeze:
|
||||
self.freeze()
|
||||
self.layer = layer
|
||||
if self.layer == 'last':
|
||||
self.layer_idx = 0
|
||||
elif self.layer == 'penultimate':
|
||||
self.layer_idx = 1
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
def freeze(self):
|
||||
self.model = self.model.eval()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, image):
|
||||
z = self.model.encode_image(image.to(self.device))
|
||||
return z
|
||||
|
||||
def encode_with_transformer(self, text):
|
||||
x = self.model.token_embedding(text)
|
||||
x = x + self.model.positional_embedding
|
||||
x = x.permute(1, 0, 2)
|
||||
x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
|
||||
x = x.permute(1, 0, 2)
|
||||
x = self.model.ln_final(x)
|
||||
|
||||
return x
|
||||
|
||||
def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
|
||||
for i, r in enumerate(self.model.transformer.resblocks):
|
||||
if i == len(self.model.transformer.resblocks) - self.layer_idx:
|
||||
break
|
||||
if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(
|
||||
):
|
||||
x = checkpoint(r, x, attn_mask)
|
||||
else:
|
||||
x = r(x, attn_mask=attn_mask)
|
||||
return x
|
||||
|
||||
def encode(self, text):
|
||||
return self(text)
|
||||
1504
modelscope/models/multi_modal/image_to_video/modules/unet_i2v.py
Normal file
1504
modelscope/models/multi_modal/image_to_video/modules/unet_i2v.py
Normal file
File diff suppressed because it is too large
Load Diff
2
modelscope/models/multi_modal/image_to_video/utils/__init__.py
Executable file
2
modelscope/models/multi_modal/image_to_video/utils/__init__.py
Executable file
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import os
|
||||
161
modelscope/models/multi_modal/image_to_video/utils/config.py
Executable file
161
modelscope/models/multi_modal/image_to_video/utils/config.py
Executable file
@@ -0,0 +1,161 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import os.path as osp
|
||||
from datetime import datetime
|
||||
|
||||
import torch
|
||||
from easydict import EasyDict
|
||||
|
||||
cfg = EasyDict(__name__='Config: VideoLDM Decoder')
|
||||
|
||||
# ---------------------------work dir--------------------------
|
||||
cfg.work_dir = 'workspace/'
|
||||
|
||||
# ---------------------------Global Variable-----------------------------------
|
||||
cfg.resolution = [448, 256]
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Dataset Parameter---------------------------------
|
||||
cfg.mean = [0.5, 0.5, 0.5]
|
||||
cfg.std = [0.5, 0.5, 0.5]
|
||||
cfg.max_words = 1000
|
||||
|
||||
# PlaceHolder
|
||||
cfg.vit_out_dim = 1024
|
||||
cfg.vit_resolution = [224, 224]
|
||||
cfg.depth_clamp = 10.0
|
||||
cfg.misc_size = 384
|
||||
cfg.depth_std = 20.0
|
||||
|
||||
cfg.frame_lens = 32
|
||||
cfg.sample_fps = 8
|
||||
|
||||
cfg.batch_sizes = 1
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Mode Parameters-----------------------------------
|
||||
# Diffusion
|
||||
cfg.schedule = 'cosine'
|
||||
cfg.num_timesteps = 1000
|
||||
cfg.mean_type = 'v'
|
||||
cfg.var_type = 'fixed_small'
|
||||
cfg.loss_type = 'mse'
|
||||
cfg.ddim_timesteps = 50
|
||||
cfg.ddim_eta = 0.0
|
||||
cfg.clamp = 1.0
|
||||
cfg.share_noise = False
|
||||
cfg.use_div_loss = False
|
||||
cfg.noise_strength = 0.1
|
||||
|
||||
# classifier-free guidance
|
||||
cfg.p_zero = 0.1
|
||||
cfg.guide_scale = 3.0
|
||||
|
||||
# clip vision encoder
|
||||
cfg.vit_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||
cfg.vit_std = [0.26862954, 0.26130258, 0.27577711]
|
||||
|
||||
# Model
|
||||
cfg.scale_factor = 0.18215
|
||||
cfg.use_fp16 = True
|
||||
cfg.temporal_attention = True
|
||||
cfg.decoder_bs = 8
|
||||
|
||||
cfg.UNet = {
|
||||
'type': 'Img2VidSDUNet',
|
||||
'in_dim': 4,
|
||||
'dim': 320,
|
||||
'y_dim': cfg.vit_out_dim,
|
||||
'context_dim': 1024,
|
||||
'out_dim': 8 if cfg.var_type.startswith('learned') else 4,
|
||||
'dim_mult': [1, 2, 4, 4],
|
||||
'num_heads': 8,
|
||||
'head_dim': 64,
|
||||
'num_res_blocks': 2,
|
||||
'attn_scales': [1 / 1, 1 / 2, 1 / 4],
|
||||
'dropout': 0.1,
|
||||
'temporal_attention': cfg.temporal_attention,
|
||||
'temporal_attn_times': 1,
|
||||
'use_checkpoint': False,
|
||||
'use_fps_condition': False,
|
||||
'use_sim_mask': False,
|
||||
'num_tokens': 4,
|
||||
'default_fps': 8,
|
||||
'input_dim': 1024
|
||||
}
|
||||
|
||||
cfg.guidances = []
|
||||
|
||||
# auotoencoder from stabel diffusion
|
||||
cfg.auto_encoder = {
|
||||
'type': 'AutoencoderKL',
|
||||
'ddconfig': {
|
||||
'double_z': True,
|
||||
'z_channels': 4,
|
||||
'resolution': 256,
|
||||
'in_channels': 3,
|
||||
'out_ch': 3,
|
||||
'ch': 128,
|
||||
'ch_mult': [1, 2, 4, 4],
|
||||
'num_res_blocks': 2,
|
||||
'attn_resolutions': [],
|
||||
'dropout': 0.0
|
||||
},
|
||||
'embed_dim': 4,
|
||||
'pretrained': 'v2-1_512-ema-pruned.ckpt'
|
||||
}
|
||||
# clip embedder
|
||||
cfg.embedder = {
|
||||
'type': 'FrozenOpenCLIPVisualEmbedder',
|
||||
'layer': 'penultimate',
|
||||
'vit_resolution': [224, 224],
|
||||
'pretrained': 'open_clip_pytorch_model.bin'
|
||||
}
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Training Settings---------------------------------
|
||||
# training and optimizer
|
||||
cfg.ema_decay = 0.9999
|
||||
cfg.num_steps = 600000
|
||||
cfg.lr = 5e-5
|
||||
cfg.weight_decay = 0.0
|
||||
cfg.betas = (0.9, 0.999)
|
||||
cfg.eps = 1.0e-8
|
||||
cfg.chunk_size = 16
|
||||
cfg.alpha = 0.7
|
||||
cfg.save_ckp_interval = 1000
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ----------------------------Pretrain Settings---------------------------------
|
||||
# Default: load 2d pretrain
|
||||
cfg.fix_weight = False
|
||||
cfg.load_match = False
|
||||
cfg.pretrained_checkpoint = 'v2-1_512-ema-pruned.ckpt'
|
||||
cfg.pretrained_image_keys = 'stable_diffusion_image_key_temporal_attention_x1.json'
|
||||
cfg.resume_checkpoint = 'img2video_ldm_0779000.pth'
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# -----------------------------Visual-------------------------------------------
|
||||
# Visual videos
|
||||
cfg.viz_interval = 1000
|
||||
cfg.visual_train = {
|
||||
'type': 'VisualVideoTextDuringTrain',
|
||||
}
|
||||
cfg.visual_inference = {
|
||||
'type': 'VisualGeneratedVideos',
|
||||
}
|
||||
cfg.inference_list_path = ''
|
||||
|
||||
# logging
|
||||
cfg.log_interval = 100
|
||||
|
||||
# Default log_dir
|
||||
cfg.log_dir = 'workspace/output_data'
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Others--------------------------------------------
|
||||
# seed
|
||||
cfg.seed = 8888
|
||||
# -----------------------------------------------------------------------------
|
||||
511
modelscope/models/multi_modal/image_to_video/utils/diffusion.py
Executable file
511
modelscope/models/multi_modal/image_to_video/utils/diffusion.py
Executable file
@@ -0,0 +1,511 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ['GaussianDiffusion', 'beta_schedule']
|
||||
|
||||
|
||||
def _i(tensor, t, x):
|
||||
r"""Index tensor using t and format the output according to x.
|
||||
"""
|
||||
shape = (x.size(0), ) + (1, ) * (x.ndim - 1)
|
||||
if tensor.device != x.device:
|
||||
tensor = tensor.to(x.device)
|
||||
return tensor[t].view(shape).to(x)
|
||||
|
||||
|
||||
def fn(u):
|
||||
return math.cos((u + 0.008) / 1.008 * math.pi / 2)**2
|
||||
|
||||
|
||||
def beta_schedule(schedule,
|
||||
num_timesteps=1000,
|
||||
init_beta=None,
|
||||
last_beta=None):
|
||||
if schedule == 'linear':
|
||||
scale = 1000.0 / num_timesteps
|
||||
init_beta = init_beta or scale * 0.0001
|
||||
last_beta = last_beta or scale * 0.02
|
||||
return torch.linspace(
|
||||
init_beta, last_beta, num_timesteps, dtype=torch.float64)
|
||||
elif schedule == 'quadratic':
|
||||
init_beta = init_beta or 0.0015
|
||||
last_beta = last_beta or 0.0195
|
||||
return torch.linspace(
|
||||
init_beta**0.5, last_beta**0.5, num_timesteps,
|
||||
dtype=torch.float64)**2
|
||||
elif schedule == 'cosine':
|
||||
betas = []
|
||||
for step in range(num_timesteps):
|
||||
t1 = step / num_timesteps
|
||||
t2 = (step + 1) / num_timesteps
|
||||
betas.append(min(1.0 - fn(t2) / fn(t1), 0.999))
|
||||
return torch.tensor(betas, dtype=torch.float64)
|
||||
else:
|
||||
raise ValueError(f'Unsupported schedule: {schedule}')
|
||||
|
||||
|
||||
class GaussianDiffusion(object):
|
||||
|
||||
def __init__(self,
|
||||
betas,
|
||||
mean_type='eps',
|
||||
var_type='learned_range',
|
||||
loss_type='mse',
|
||||
epsilon=1e-12,
|
||||
rescale_timesteps=False,
|
||||
noise_strength=0.0):
|
||||
# check input
|
||||
if not isinstance(betas, torch.DoubleTensor):
|
||||
betas = torch.tensor(betas, dtype=torch.float64)
|
||||
assert min(betas) > 0 and max(betas) <= 1
|
||||
assert mean_type in ['x0', 'x_{t-1}', 'eps', 'v']
|
||||
assert var_type in [
|
||||
'learned', 'learned_range', 'fixed_large', 'fixed_small'
|
||||
]
|
||||
assert loss_type in [
|
||||
'mse', 'rescaled_mse', 'kl', 'rescaled_kl', 'l1', 'rescaled_l1',
|
||||
'charbonnier'
|
||||
]
|
||||
self.betas = betas
|
||||
self.num_timesteps = len(betas)
|
||||
self.mean_type = mean_type
|
||||
self.var_type = var_type
|
||||
self.loss_type = loss_type
|
||||
self.epsilon = epsilon
|
||||
self.rescale_timesteps = rescale_timesteps
|
||||
self.noise_strength = noise_strength
|
||||
|
||||
# alphas
|
||||
alphas = 1 - self.betas
|
||||
self.alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||||
self.alphas_cumprod_prev = torch.cat(
|
||||
[alphas.new_ones([1]), self.alphas_cumprod[:-1]])
|
||||
self.alphas_cumprod_next = torch.cat(
|
||||
[self.alphas_cumprod[1:],
|
||||
alphas.new_zeros([1])])
|
||||
|
||||
# q(x_t | x_{t-1})
|
||||
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
|
||||
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0
|
||||
- self.alphas_cumprod)
|
||||
self.log_one_minus_alphas_cumprod = torch.log(1.0
|
||||
- self.alphas_cumprod)
|
||||
self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)
|
||||
self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod
|
||||
- 1)
|
||||
|
||||
# q(x_{t-1} | x_t, x_0)
|
||||
self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (
|
||||
1.0 - self.alphas_cumprod)
|
||||
self.posterior_log_variance_clipped = torch.log(
|
||||
self.posterior_variance.clamp(1e-20))
|
||||
self.posterior_mean_coef1 = betas * torch.sqrt(
|
||||
self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
|
||||
self.posterior_mean_coef2 = (
|
||||
1.0 - self.alphas_cumprod_prev) * torch.sqrt(alphas) / (
|
||||
1.0 - self.alphas_cumprod)
|
||||
|
||||
def sample_loss(self, x0, noise=None):
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x0)
|
||||
if self.noise_strength > 0:
|
||||
b, c, f, _, _ = x0.shape
|
||||
offset_noise = torch.randn(b, c, f, 1, 1, device=x0.device)
|
||||
noise = noise + self.noise_strength * offset_noise
|
||||
return noise
|
||||
|
||||
def q_sample(self, x0, t, noise=None):
|
||||
r"""Sample from q(x_t | x_0).
|
||||
"""
|
||||
# noise = torch.randn_like(x0) if noise is None else noise
|
||||
noise = self.sample_loss(x0, noise)
|
||||
return _i(self.sqrt_alphas_cumprod, t, x0) * x0 + (
|
||||
_i(self.sqrt_one_minus_alphas_cumprod, t, x0) * noise)
|
||||
|
||||
def q_mean_variance(self, x0, t):
|
||||
r"""Distribution of q(x_t | x_0).
|
||||
"""
|
||||
mu = _i(self.sqrt_alphas_cumprod, t, x0) * x0
|
||||
var = _i(1.0 - self.alphas_cumprod, t, x0)
|
||||
log_var = _i(self.log_one_minus_alphas_cumprod, t, x0)
|
||||
return mu, var, log_var
|
||||
|
||||
def q_posterior_mean_variance(self, x0, xt, t):
|
||||
r"""Distribution of q(x_{t-1} | x_t, x_0).
|
||||
"""
|
||||
mu = _i(self.posterior_mean_coef1, t, xt) * x0 + _i(
|
||||
self.posterior_mean_coef2, t, xt) * xt
|
||||
var = _i(self.posterior_variance, t, xt)
|
||||
log_var = _i(self.posterior_log_variance_clipped, t, xt)
|
||||
return mu, var, log_var
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample(self,
|
||||
xt,
|
||||
t,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None):
|
||||
r"""Sample from p(x_{t-1} | x_t).
|
||||
- condition_fn: for classifier-based guidance (guided-diffusion).
|
||||
- guide_scale: for classifier-free guidance (glide/dalle-2).
|
||||
"""
|
||||
# predict distribution of p(x_{t-1} | x_t)
|
||||
mu, var, log_var, x0 = self.p_mean_variance(xt, t, model, model_kwargs,
|
||||
clamp, percentile,
|
||||
guide_scale)
|
||||
|
||||
# random sample (with optional conditional function)
|
||||
noise = torch.randn_like(xt)
|
||||
mask = t.ne(0).float().view(-1, *((1, ) * (xt.ndim - 1)))
|
||||
if condition_fn is not None:
|
||||
grad = condition_fn(xt, self._scale_timesteps(t), **model_kwargs)
|
||||
mu = mu.float() + var * grad.float()
|
||||
xt_1 = mu + mask * torch.exp(0.5 * log_var) * noise
|
||||
return xt_1, x0
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample_loop(self,
|
||||
noise,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None):
|
||||
r"""Sample from p(x_{t-1} | x_t) p(x_{t-2} | x_{t-1}) ... p(x_0 | x_1).
|
||||
"""
|
||||
# prepare input
|
||||
b = noise.size(0)
|
||||
xt = noise
|
||||
|
||||
# diffusion process
|
||||
for step in torch.arange(self.num_timesteps).flip(0):
|
||||
t = torch.full((b, ), step, dtype=torch.long, device=xt.device)
|
||||
xt, _ = self.p_sample(xt, t, model, model_kwargs, clamp,
|
||||
percentile, condition_fn, guide_scale)
|
||||
return xt
|
||||
|
||||
def p_mean_variance(self,
|
||||
xt,
|
||||
t,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
guide_scale=None):
|
||||
r"""Distribution of p(x_{t-1} | x_t).
|
||||
"""
|
||||
# predict distribution
|
||||
if guide_scale is None:
|
||||
out = model(xt, self._scale_timesteps(t), **model_kwargs)
|
||||
else:
|
||||
# classifier-free guidance
|
||||
# (model_kwargs[0]: conditional kwargs; model_kwargs[1]: non-conditional kwargs)
|
||||
assert isinstance(model_kwargs, list) and len(model_kwargs) == 2
|
||||
y_out = model(xt, self._scale_timesteps(t), **model_kwargs[0])
|
||||
u_out = model(xt, self._scale_timesteps(t), **model_kwargs[1])
|
||||
dim = y_out.size(1) if self.var_type.startswith(
|
||||
'fixed') else y_out.size(1) // 2
|
||||
out = torch.cat(
|
||||
[
|
||||
u_out[:, :dim] + guide_scale * # noqa
|
||||
(y_out[:, :dim] - u_out[:, :dim]),
|
||||
y_out[:, dim:]
|
||||
],
|
||||
dim=1)
|
||||
|
||||
# compute variance
|
||||
if self.var_type == 'learned':
|
||||
out, log_var = out.chunk(2, dim=1)
|
||||
var = torch.exp(log_var)
|
||||
elif self.var_type == 'learned_range':
|
||||
out, fraction = out.chunk(2, dim=1)
|
||||
min_log_var = _i(self.posterior_log_variance_clipped, t, xt)
|
||||
max_log_var = _i(torch.log(self.betas), t, xt)
|
||||
fraction = (fraction + 1) / 2.0
|
||||
log_var = fraction * max_log_var + (1 - fraction) * min_log_var
|
||||
var = torch.exp(log_var)
|
||||
elif self.var_type == 'fixed_large':
|
||||
var = _i(
|
||||
torch.cat([self.posterior_variance[1:2], self.betas[1:]]), t,
|
||||
xt)
|
||||
log_var = torch.log(var)
|
||||
elif self.var_type == 'fixed_small':
|
||||
var = _i(self.posterior_variance, t, xt)
|
||||
log_var = _i(self.posterior_log_variance_clipped, t, xt)
|
||||
|
||||
# compute mean and x0
|
||||
if self.mean_type == 'x_{t-1}':
|
||||
mu = out
|
||||
x0 = _i(1.0 / self.posterior_mean_coef1, t, xt) * mu - (
|
||||
_i(self.posterior_mean_coef2 / self.posterior_mean_coef1, t,
|
||||
xt) * xt)
|
||||
elif self.mean_type == 'x0':
|
||||
x0 = out
|
||||
mu, _, _ = self.q_posterior_mean_variance(x0, xt, t)
|
||||
elif self.mean_type == 'eps':
|
||||
x0 = _i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt) * out)
|
||||
mu, _, _ = self.q_posterior_mean_variance(x0, xt, t)
|
||||
elif self.mean_type == 'v':
|
||||
x0 = _i(self.sqrt_alphas_cumprod, t, xt) * xt - (
|
||||
_i(self.sqrt_one_minus_alphas_cumprod, t, xt) * out)
|
||||
mu, _, _ = self.q_posterior_mean_variance(x0, xt, t)
|
||||
|
||||
# restrict the range of x0
|
||||
if percentile is not None:
|
||||
assert percentile > 0 and percentile <= 1
|
||||
s = torch.quantile(
|
||||
x0.flatten(1).abs(), percentile,
|
||||
dim=1).clamp_(1.0).view(-1, 1, 1, 1)
|
||||
x0 = torch.min(s, torch.max(-s, x0)) / s
|
||||
elif clamp is not None:
|
||||
x0 = x0.clamp(-clamp, clamp)
|
||||
return mu, var, log_var, x0
|
||||
|
||||
@torch.no_grad()
|
||||
def ddim_sample(self,
|
||||
xt,
|
||||
t,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
ddim_timesteps=20,
|
||||
eta=0.0):
|
||||
r"""Sample from p(x_{t-1} | x_t) using DDIM.
|
||||
- condition_fn: for classifier-based guidance (guided-diffusion).
|
||||
- guide_scale: for classifier-free guidance (glide/dalle-2).
|
||||
"""
|
||||
stride = self.num_timesteps // ddim_timesteps
|
||||
|
||||
# predict distribution of p(x_{t-1} | x_t)
|
||||
_, _, _, x0 = self.p_mean_variance(xt, t, model, model_kwargs, clamp,
|
||||
percentile, guide_scale)
|
||||
if condition_fn is not None:
|
||||
# x0 -> eps
|
||||
alpha = _i(self.alphas_cumprod, t, xt)
|
||||
eps = (_i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - x0) / (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt))
|
||||
eps = eps - (1 - alpha).sqrt() * condition_fn(
|
||||
xt, self._scale_timesteps(t), **model_kwargs)
|
||||
|
||||
# eps -> x0
|
||||
x0 = _i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt) * eps)
|
||||
|
||||
# derive variables
|
||||
eps = (_i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - x0) / (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt))
|
||||
alphas = _i(self.alphas_cumprod, t, xt)
|
||||
alphas_prev = _i(self.alphas_cumprod, (t - stride).clamp(0), xt)
|
||||
sigmas = eta * torch.sqrt((1 - alphas_prev) / (1 - alphas) * # noqa
|
||||
(1 - alphas / alphas_prev))
|
||||
|
||||
# random sample
|
||||
noise = torch.randn_like(xt)
|
||||
direction = torch.sqrt(1 - alphas_prev - sigmas**2) * eps
|
||||
mask = t.ne(0).float().view(-1, *((1, ) * (xt.ndim - 1)))
|
||||
xt_1 = torch.sqrt(alphas_prev) * x0 + direction + mask * sigmas * noise
|
||||
return xt_1, x0
|
||||
|
||||
@torch.no_grad()
|
||||
def ddim_sample_loop(self,
|
||||
noise,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
ddim_timesteps=20,
|
||||
eta=0.0):
|
||||
# prepare input
|
||||
b = noise.size(0)
|
||||
xt = noise
|
||||
|
||||
# diffusion process (TODO: clamp is inaccurate! Consider replacing the stride by explicit prev/next steps)
|
||||
steps = (1 + torch.arange(0, self.num_timesteps,
|
||||
self.num_timesteps // ddim_timesteps)).clamp(
|
||||
0, self.num_timesteps - 1).flip(0)
|
||||
for step in steps:
|
||||
t = torch.full((b, ), step, dtype=torch.long, device=xt.device)
|
||||
xt, _ = self.ddim_sample(xt, t, model, model_kwargs, clamp,
|
||||
percentile, condition_fn, guide_scale,
|
||||
ddim_timesteps, eta)
|
||||
return xt
|
||||
|
||||
@torch.no_grad()
|
||||
def ddim_reverse_sample(self,
|
||||
xt,
|
||||
t,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
guide_scale=None,
|
||||
ddim_timesteps=20):
|
||||
r"""Sample from p(x_{t+1} | x_t) using DDIM reverse ODE (deterministic).
|
||||
"""
|
||||
stride = self.num_timesteps // ddim_timesteps
|
||||
|
||||
# predict distribution of p(x_{t-1} | x_t)
|
||||
_, _, _, x0 = self.p_mean_variance(xt, t, model, model_kwargs, clamp,
|
||||
percentile, guide_scale)
|
||||
|
||||
# derive variables
|
||||
eps = (_i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - x0) / (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt))
|
||||
alphas_next = _i(
|
||||
torch.cat(
|
||||
[self.alphas_cumprod,
|
||||
self.alphas_cumprod.new_zeros([1])]),
|
||||
(t + stride).clamp(0, self.num_timesteps), xt)
|
||||
|
||||
# reverse sample
|
||||
mu = torch.sqrt(alphas_next) * x0 + torch.sqrt(1 - alphas_next) * eps
|
||||
return mu, x0
|
||||
|
||||
@torch.no_grad()
|
||||
def ddim_reverse_sample_loop(self,
|
||||
x0,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
guide_scale=None,
|
||||
ddim_timesteps=20):
|
||||
# prepare input
|
||||
b = x0.size(0)
|
||||
xt = x0
|
||||
|
||||
# reconstruction steps
|
||||
steps = torch.arange(0, self.num_timesteps,
|
||||
self.num_timesteps // ddim_timesteps)
|
||||
for step in steps:
|
||||
t = torch.full((b, ), step, dtype=torch.long, device=xt.device)
|
||||
xt, _ = self.ddim_reverse_sample(xt, t, model, model_kwargs, clamp,
|
||||
percentile, guide_scale,
|
||||
ddim_timesteps)
|
||||
return xt
|
||||
|
||||
@torch.no_grad()
|
||||
def plms_sample(self,
|
||||
xt,
|
||||
t,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
plms_timesteps=20):
|
||||
r"""Sample from p(x_{t-1} | x_t) using PLMS.
|
||||
- condition_fn: for classifier-based guidance (guided-diffusion).
|
||||
- guide_scale: for classifier-free guidance (glide/dalle-2).
|
||||
"""
|
||||
stride = self.num_timesteps // plms_timesteps
|
||||
|
||||
# function for compute eps
|
||||
def compute_eps(xt, t):
|
||||
# predict distribution of p(x_{t-1} | x_t)
|
||||
_, _, _, x0 = self.p_mean_variance(xt, t, model, model_kwargs,
|
||||
clamp, percentile, guide_scale)
|
||||
|
||||
# condition
|
||||
if condition_fn is not None:
|
||||
# x0 -> eps
|
||||
alpha = _i(self.alphas_cumprod, t, xt)
|
||||
eps = (_i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - x0) / (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt))
|
||||
eps = eps - (1 - alpha).sqrt() * condition_fn(
|
||||
xt, self._scale_timesteps(t), **model_kwargs)
|
||||
|
||||
# eps -> x0
|
||||
x0 = _i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt) * eps)
|
||||
|
||||
# derive eps
|
||||
eps = (_i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - x0) / (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt))
|
||||
return eps
|
||||
|
||||
# function for compute x_0 and x_{t-1}
|
||||
def compute_x0(eps, t):
|
||||
# eps -> x0
|
||||
x0 = _i(self.sqrt_recip_alphas_cumprod, t, xt) * xt - (
|
||||
_i(self.sqrt_recipm1_alphas_cumprod, t, xt) * eps)
|
||||
|
||||
# deterministic sample
|
||||
alphas_prev = _i(self.alphas_cumprod, (t - stride).clamp(0), xt)
|
||||
direction = torch.sqrt(1 - alphas_prev) * eps
|
||||
xt_1 = torch.sqrt(alphas_prev) * x0 + direction
|
||||
return xt_1, x0
|
||||
|
||||
# PLMS sample
|
||||
eps = compute_eps(xt, t)
|
||||
if len(eps_cache) == 0:
|
||||
# 2nd order pseudo improved Euler
|
||||
xt_1, x0 = compute_x0(eps, t)
|
||||
eps_next = compute_eps(xt_1, (t - stride).clamp(0))
|
||||
eps_prime = (eps + eps_next) / 2.0
|
||||
elif len(eps_cache) == 1:
|
||||
# 2nd order pseudo linear multistep (Adams-Bashforth)
|
||||
eps_prime = (3 * eps - eps_cache[-1]) / 2.0
|
||||
elif len(eps_cache) == 2:
|
||||
# 3nd order pseudo linear multistep (Adams-Bashforth)
|
||||
eps_prime = (23 * eps - 16 * eps_cache[-1]
|
||||
+ 5 * eps_cache[-2]) / 12.0
|
||||
elif len(eps_cache) >= 3:
|
||||
# 4nd order pseudo linear multistep (Adams-Bashforth)
|
||||
eps_prime = (55 * eps - 59 * eps_cache[-1] + 37 * eps_cache[-2]
|
||||
- 9 * eps_cache[-3]) / 24.0
|
||||
xt_1, x0 = compute_x0(eps_prime, t)
|
||||
return xt_1, x0, eps
|
||||
|
||||
@torch.no_grad()
|
||||
def plms_sample_loop(self,
|
||||
noise,
|
||||
model,
|
||||
model_kwargs={},
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
plms_timesteps=20):
|
||||
# prepare input
|
||||
b = noise.size(0)
|
||||
xt = noise
|
||||
|
||||
# diffusion process
|
||||
steps = (1 + torch.arange(0, self.num_timesteps,
|
||||
self.num_timesteps // plms_timesteps)).clamp(
|
||||
0, self.num_timesteps - 1).flip(0)
|
||||
eps_cache = []
|
||||
for step in steps:
|
||||
# PLMS sampling step
|
||||
t = torch.full((b, ), step, dtype=torch.long, device=xt.device)
|
||||
xt, _, eps = self.plms_sample(xt, t, model, model_kwargs, clamp,
|
||||
percentile, condition_fn,
|
||||
guide_scale, plms_timesteps,
|
||||
eps_cache)
|
||||
|
||||
# update eps cache
|
||||
eps_cache.append(eps)
|
||||
if len(eps_cache) >= 4:
|
||||
eps_cache.pop(0)
|
||||
return xt
|
||||
|
||||
def _scale_timesteps(self, t):
|
||||
if self.rescale_timesteps:
|
||||
return t.float() * 1000.0 / self.num_timesteps
|
||||
return t
|
||||
14
modelscope/models/multi_modal/image_to_video/utils/seed.py
Executable file
14
modelscope/models/multi_modal/image_to_video/utils/seed.py
Executable file
@@ -0,0 +1,14 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def setup_seed(seed):
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def fn(u):
|
||||
return math.cos((u + 0.008) / 1.008 * math.pi / 2)**2
|
||||
|
||||
|
||||
def beta_schedule(schedule,
|
||||
num_timesteps=1000,
|
||||
init_beta=None,
|
||||
last_beta=None):
|
||||
'''
|
||||
This code defines a function beta_schedule that generates a sequence of beta values based on the given input
|
||||
parameters. These beta values can be used in video diffusion processes. The function has the following parameters:
|
||||
schedule(str): Determines the type of beta schedule to be generated. It can be 'linear', 'linear_sd',
|
||||
'quadratic', or 'cosine'.
|
||||
num_timesteps(int, optional): The number of timesteps for the generated beta schedule. Default is 1000.
|
||||
init_beta(float, optional): The initial beta value. If not provided, a default value is used based on the
|
||||
chosen schedule.
|
||||
last_beta(float, optional): The final beta value. If not provided, a default value is used based on the
|
||||
chosen schedule.
|
||||
The function returns a PyTorch tensor containing the generated beta values.
|
||||
The beta schedule is determined by the schedule parameter:
|
||||
1.Linear: Generates a linear sequence of beta values betweeninit_betaandlast_beta.
|
||||
2.Linear_sd: Generates a linear sequence of beta values between the square root of init_beta and the square root
|
||||
oflast_beta, and then squares the result.
|
||||
3.Quadratic: Similar to the 'linear_sd' schedule, but with different default values forinit_betaandlast_beta.
|
||||
4.Cosine: Generates a sequence of beta values based on a cosine function, ensuring the values are between 0
|
||||
and 0.999.
|
||||
If an unsupported schedule is provided, a ValueError is raised with a message indicating the issue.
|
||||
'''
|
||||
if schedule == 'linear':
|
||||
scale = 1000.0 / num_timesteps
|
||||
init_beta = init_beta or scale * 0.0001
|
||||
last_beta = last_beta or scale * 0.02
|
||||
return torch.linspace(
|
||||
init_beta, last_beta, num_timesteps, dtype=torch.float64)
|
||||
elif schedule == 'linear_sd':
|
||||
return torch.linspace(
|
||||
init_beta**0.5, last_beta**0.5, num_timesteps,
|
||||
dtype=torch.float64)**2
|
||||
elif schedule == 'quadratic':
|
||||
init_beta = init_beta or 0.0015
|
||||
last_beta = last_beta or 0.0195
|
||||
return torch.linspace(
|
||||
init_beta**0.5, last_beta**0.5, num_timesteps,
|
||||
dtype=torch.float64)**2
|
||||
elif schedule == 'cosine':
|
||||
betas = []
|
||||
for step in range(num_timesteps):
|
||||
t1 = step / num_timesteps
|
||||
t2 = (step + 1) / num_timesteps
|
||||
betas.append(min(1.0 - fn(t2) / fn(t1), 0.999))
|
||||
return torch.tensor(betas, dtype=torch.float64)
|
||||
else:
|
||||
raise ValueError(f'Unsupported schedule: {schedule}')
|
||||
404
modelscope/models/multi_modal/image_to_video/utils/transforms.py
Executable file
404
modelscope/models/multi_modal/image_to_video/utils/transforms.py
Executable file
@@ -0,0 +1,404 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as F
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
__all__ = [
|
||||
'Compose', 'Resize', 'Rescale', 'CenterCrop', 'CenterCropV2',
|
||||
'CenterCropWide', 'RandomCrop', 'RandomCropV2', 'RandomHFlip',
|
||||
'GaussianBlur', 'ColorJitter', 'RandomGray', 'ToTensor', 'Normalize',
|
||||
'ResizeRandomCrop', 'ExtractResizeRandomCrop', 'ExtractResizeAssignCrop'
|
||||
]
|
||||
|
||||
|
||||
class Compose(object):
|
||||
|
||||
def __init__(self, transforms):
|
||||
self.transforms = transforms
|
||||
|
||||
def __getitem__(self, index):
|
||||
if isinstance(index, slice):
|
||||
return Compose(self.transforms[index])
|
||||
else:
|
||||
return self.transforms[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.transforms)
|
||||
|
||||
def __call__(self, rgb):
|
||||
for t in self.transforms:
|
||||
rgb = t(rgb)
|
||||
return rgb
|
||||
|
||||
|
||||
class Resize(object):
|
||||
|
||||
def __init__(self, size=256):
|
||||
if isinstance(size, int):
|
||||
size = (size, size)
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
if isinstance(rgb, list):
|
||||
rgb = [u.resize(self.size, Image.BILINEAR) for u in rgb]
|
||||
else:
|
||||
rgb = rgb.resize(self.size, Image.BILINEAR)
|
||||
return rgb
|
||||
|
||||
|
||||
class Rescale(object):
|
||||
|
||||
def __init__(self, size=256, interpolation=Image.BILINEAR):
|
||||
self.size = size
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
scale = self.size / min(w, h)
|
||||
out_w, out_h = int(round(w * scale)), int(round(h * scale))
|
||||
rgb = [u.resize((out_w, out_h), self.interpolation) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCrop(object):
|
||||
|
||||
def __init__(self, size=224):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
assert min(w, h) >= self.size
|
||||
x1 = (w - self.size) // 2
|
||||
y1 = (h - self.size) // 2
|
||||
rgb = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ExtractResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
wh = [x1, y1, x1 + out_w, y1 + out_h]
|
||||
return rgb, wh
|
||||
|
||||
|
||||
class ExtractResizeAssignCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb, wh):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
|
||||
rgb = [u.crop(wh) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCropV2(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
# fast resize
|
||||
while min(img[0].size) >= 2 * self.size:
|
||||
img = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in img
|
||||
]
|
||||
scale = self.size / min(img[0].size)
|
||||
img = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in img
|
||||
]
|
||||
|
||||
# center crop
|
||||
x1 = (img[0].width - self.size) // 2
|
||||
y1 = (img[0].height - self.size) // 2
|
||||
img = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in img]
|
||||
return img
|
||||
|
||||
|
||||
class CenterCropWide(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
if isinstance(img, list):
|
||||
scale = min(img[0].size[0] / self.size[0],
|
||||
img[0].size[1] / self.size[1])
|
||||
img = [
|
||||
u.resize((round(u.width // scale), round(u.height // scale)),
|
||||
resample=Image.BOX) for u in img
|
||||
]
|
||||
|
||||
# center crop
|
||||
x1 = (img[0].width - self.size[0]) // 2
|
||||
y1 = (img[0].height - self.size[1]) // 2
|
||||
img = [
|
||||
u.crop((x1, y1, x1 + self.size[0], y1 + self.size[1]))
|
||||
for u in img
|
||||
]
|
||||
return img
|
||||
else:
|
||||
scale = min(img.size[0] / self.size[0], img.size[1] / self.size[1])
|
||||
img = img.resize(
|
||||
(round(img.width // scale), round(img.height // scale)),
|
||||
resample=Image.BOX)
|
||||
x1 = (img.width - self.size[0]) // 2
|
||||
y1 = (img.height - self.size[1]) // 2
|
||||
img = img.crop((x1, y1, x1 + self.size[0], y1 + self.size[1]))
|
||||
return img
|
||||
|
||||
|
||||
class RandomCrop(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4):
|
||||
self.size = size
|
||||
self.min_area = min_area
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
w, h = rgb[0].size
|
||||
area = w * h
|
||||
out_w, out_h = float('inf'), float('inf')
|
||||
while out_w > w or out_h > h:
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
aspect_ratio = random.uniform(3. / 4., 4. / 3.)
|
||||
out_w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
out_h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomCropV2(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4, ratio=(3. / 4., 4. / 3.)):
|
||||
if isinstance(size, (tuple, list)):
|
||||
self.size = size
|
||||
else:
|
||||
self.size = (size, size)
|
||||
self.min_area = min_area
|
||||
self.ratio = ratio
|
||||
|
||||
def _get_params(self, img):
|
||||
width, height = img.size
|
||||
area = height * width
|
||||
|
||||
for _ in range(10):
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
log_ratio = (math.log(self.ratio[0]), math.log(self.ratio[1]))
|
||||
aspect_ratio = math.exp(random.uniform(*log_ratio))
|
||||
|
||||
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
|
||||
if 0 < w <= width and 0 < h <= height:
|
||||
i = random.randint(0, height - h)
|
||||
j = random.randint(0, width - w)
|
||||
return i, j, h, w
|
||||
|
||||
# Fallback to central crop
|
||||
in_ratio = float(width) / float(height)
|
||||
if (in_ratio < min(self.ratio)):
|
||||
w = width
|
||||
h = int(round(w / min(self.ratio)))
|
||||
elif (in_ratio > max(self.ratio)):
|
||||
h = height
|
||||
w = int(round(h * max(self.ratio)))
|
||||
else:
|
||||
w = width
|
||||
h = height
|
||||
i = (height - h) // 2
|
||||
j = (width - w) // 2
|
||||
return i, j, h, w
|
||||
|
||||
def __call__(self, rgb):
|
||||
i, j, h, w = self._get_params(rgb[0])
|
||||
rgb = [F.resized_crop(u, i, j, h, w, self.size) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomHFlip(object):
|
||||
|
||||
def __init__(self, p=0.5):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.transpose(Image.FLIP_LEFT_RIGHT) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class GaussianBlur(object):
|
||||
|
||||
def __init__(self, sigmas=[0.1, 2.0], p=0.5):
|
||||
self.sigmas = sigmas
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
sigma = random.uniform(*self.sigmas)
|
||||
rgb = [
|
||||
u.filter(ImageFilter.GaussianBlur(radius=sigma)) for u in rgb
|
||||
]
|
||||
return rgb
|
||||
|
||||
|
||||
class ColorJitter(object):
|
||||
|
||||
def __init__(self,
|
||||
brightness=0.4,
|
||||
contrast=0.4,
|
||||
saturation=0.4,
|
||||
hue=0.1,
|
||||
p=0.5):
|
||||
self.brightness = brightness
|
||||
self.contrast = contrast
|
||||
self.saturation = saturation
|
||||
self.hue = hue
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
brightness, contrast, saturation, hue = self._random_params()
|
||||
transforms = [
|
||||
lambda f: F.adjust_brightness(f, brightness),
|
||||
lambda f: F.adjust_contrast(f, contrast),
|
||||
lambda f: F.adjust_saturation(f, saturation),
|
||||
lambda f: F.adjust_hue(f, hue)
|
||||
]
|
||||
random.shuffle(transforms)
|
||||
for t in transforms:
|
||||
rgb = [t(u) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
def _random_params(self):
|
||||
brightness = random.uniform(
|
||||
max(0, 1 - self.brightness), 1 + self.brightness)
|
||||
contrast = random.uniform(max(0, 1 - self.contrast), 1 + self.contrast)
|
||||
saturation = random.uniform(
|
||||
max(0, 1 - self.saturation), 1 + self.saturation)
|
||||
hue = random.uniform(-self.hue, self.hue)
|
||||
return brightness, contrast, saturation, hue
|
||||
|
||||
|
||||
class RandomGray(object):
|
||||
|
||||
def __init__(self, p=0.2):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.convert('L').convert('RGB') for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
|
||||
def __call__(self, rgb):
|
||||
if isinstance(rgb, list):
|
||||
rgb = torch.stack([F.to_tensor(u) for u in rgb], dim=0)
|
||||
else:
|
||||
rgb = F.to_tensor(rgb)
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
|
||||
def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
def __call__(self, rgb):
|
||||
rgb = rgb.clone()
|
||||
rgb.clamp_(0, 1)
|
||||
if not isinstance(self.mean, torch.Tensor):
|
||||
self.mean = rgb.new_tensor(self.mean).view(-1)
|
||||
if not isinstance(self.std, torch.Tensor):
|
||||
self.std = rgb.new_tensor(self.std).view(-1)
|
||||
if rgb.dim() == 4:
|
||||
rgb.sub_(self.mean.view(1, -1, 1,
|
||||
1)).div_(self.std.view(1, -1, 1, 1))
|
||||
elif rgb.dim() == 3:
|
||||
rgb.sub_(self.mean.view(-1, 1, 1)).div_(self.std.view(-1, 1, 1))
|
||||
return rgb
|
||||
@@ -1,2 +1,3 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from .stable_diffusion import StableDiffusion
|
||||
from .stable_diffusion_xl import StableDiffusionXL
|
||||
|
||||
@@ -20,7 +20,7 @@ from modelscope.utils.constant import Tasks
|
||||
@MODELS.register_module(
|
||||
Tasks.text_to_image_synthesis, module_name=Models.stable_diffusion)
|
||||
class StableDiffusion(TorchModel):
|
||||
""" The implementation of efficient diffusion tuning model based on TorchModel.
|
||||
""" The implementation of stable diffusion model based on TorchModel.
|
||||
|
||||
This model is constructed with the implementation of stable diffusion model. If you want to
|
||||
finetune lightweight parameters on your own dataset, you can define you own tuner module
|
||||
@@ -28,7 +28,7 @@ class StableDiffusion(TorchModel):
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir, *args, **kwargs):
|
||||
""" Initialize a vision efficient diffusion tuning model.
|
||||
""" Initialize a vision stable diffusion model.
|
||||
|
||||
Args:
|
||||
model_dir: model id or path
|
||||
@@ -39,7 +39,7 @@ class StableDiffusion(TorchModel):
|
||||
self.lora_tune = kwargs.pop('lora_tune', False)
|
||||
self.dreambooth_tune = kwargs.pop('dreambooth_tune', False)
|
||||
|
||||
self.weight_dtype = torch.float32
|
||||
self.weight_dtype = kwargs.pop('torch_type', torch.float32)
|
||||
self.device = torch.device(
|
||||
'cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
@@ -59,14 +59,15 @@ class StableDiffusion(TorchModel):
|
||||
# Freeze gradient calculation and move to device
|
||||
if self.vae is not None:
|
||||
self.vae.requires_grad_(False)
|
||||
self.vae = self.vae.to(self.device)
|
||||
self.vae = self.vae.to(self.device, dtype=self.weight_dtype)
|
||||
if self.text_encoder is not None:
|
||||
self.text_encoder.requires_grad_(False)
|
||||
self.text_encoder = self.text_encoder.to(self.device)
|
||||
self.text_encoder = self.text_encoder.to(
|
||||
self.device, dtype=self.weight_dtype)
|
||||
if self.unet is not None:
|
||||
if self.lora_tune:
|
||||
self.unet.requires_grad_(False)
|
||||
self.unet = self.unet.to(self.device)
|
||||
self.unet = self.unet.to(self.device, dtype=self.weight_dtype)
|
||||
|
||||
# xformers accelerate memory efficient attention
|
||||
if xformers_enable:
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
# Copyright 2023-2024 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
||||
import os
|
||||
import random
|
||||
from functools import partial
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
|
||||
from packaging import version
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms.functional import crop
|
||||
from transformers import (AutoTokenizer, CLIPTextModel,
|
||||
CLIPTextModelWithProjection)
|
||||
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.models import TorchModel
|
||||
from modelscope.models.builder import MODELS
|
||||
from modelscope.outputs import OutputKeys
|
||||
from modelscope.utils.checkpoint import save_checkpoint, save_configuration
|
||||
from modelscope.utils.constant import Tasks
|
||||
|
||||
|
||||
@MODELS.register_module(
|
||||
Tasks.text_to_image_synthesis, module_name=Models.stable_diffusion_xl)
|
||||
class StableDiffusionXL(TorchModel):
|
||||
""" The implementation of stable diffusion xl model based on TorchModel.
|
||||
|
||||
This model is constructed with the implementation of stable diffusion xl model. If you want to
|
||||
finetune lightweight parameters on your own dataset, you can define you own tuner module
|
||||
and load in this cls.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir, *args, **kwargs):
|
||||
""" Initialize a vision stable diffusion xl model.
|
||||
|
||||
Args:
|
||||
model_dir: model id or path
|
||||
"""
|
||||
super().__init__(model_dir, *args, **kwargs)
|
||||
revision = kwargs.pop('revision', None)
|
||||
xformers_enable = kwargs.pop('xformers_enable', False)
|
||||
self.lora_tune = kwargs.pop('lora_tune', False)
|
||||
self.resolution = kwargs.pop('resolution', 1024)
|
||||
self.random_flip = kwargs.pop('random_flip', True)
|
||||
|
||||
self.weight_dtype = torch.float32
|
||||
self.device = torch.device(
|
||||
'cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# Load scheduler, tokenizer and models
|
||||
self.noise_scheduler = DDPMScheduler.from_pretrained(
|
||||
model_dir, subfolder='scheduler')
|
||||
self.tokenizer_one = AutoTokenizer.from_pretrained(
|
||||
model_dir,
|
||||
subfolder='tokenizer',
|
||||
revision=revision,
|
||||
use_fast=False)
|
||||
self.tokenizer_two = AutoTokenizer.from_pretrained(
|
||||
model_dir,
|
||||
subfolder='tokenizer_2',
|
||||
revision=revision,
|
||||
use_fast=False)
|
||||
self.text_encoder_one = CLIPTextModel.from_pretrained(
|
||||
model_dir, subfolder='text_encoder', revision=revision)
|
||||
self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
|
||||
model_dir, subfolder='text_encoder_2', revision=revision)
|
||||
self.vae = AutoencoderKL.from_pretrained(
|
||||
model_dir, subfolder='vae', revision=revision)
|
||||
self.unet = UNet2DConditionModel.from_pretrained(
|
||||
model_dir, subfolder='unet', revision=revision)
|
||||
self.safety_checker = None
|
||||
|
||||
# Freeze gradient calculation and move to device
|
||||
if self.vae is not None:
|
||||
self.vae.requires_grad_(False)
|
||||
self.vae = self.vae.to(self.device)
|
||||
if self.text_encoder_one is not None:
|
||||
self.text_encoder_one.requires_grad_(False)
|
||||
self.text_encoder_one = self.text_encoder_one.to(self.device)
|
||||
if self.text_encoder_two is not None:
|
||||
self.text_encoder_two.requires_grad_(False)
|
||||
self.text_encoder_two = self.text_encoder_two.to(self.device)
|
||||
if self.unet is not None:
|
||||
if self.lora_tune:
|
||||
self.unet.requires_grad_(False)
|
||||
self.unet = self.unet.to(self.device)
|
||||
|
||||
# xformers accelerate memory efficient attention
|
||||
if xformers_enable:
|
||||
import xformers
|
||||
|
||||
xformers_version = version.parse(xformers.__version__)
|
||||
if xformers_version == version.parse('0.0.16'):
|
||||
logger.warn(
|
||||
'xFormers 0.0.16 cannot be used for training in some GPUs. '
|
||||
'If you observe problems during training, please update xFormers to at least 0.0.17.'
|
||||
)
|
||||
self.unet.enable_xformers_memory_efficient_attention()
|
||||
|
||||
def tokenize_caption(self, tokenizer, captions):
|
||||
""" Convert caption text to token data.
|
||||
|
||||
Args:
|
||||
tokenizer: the tokenizer one or two.
|
||||
captions: a batch of texts.
|
||||
Returns: token's data as tensor.
|
||||
"""
|
||||
inputs = tokenizer(
|
||||
captions,
|
||||
max_length=tokenizer.model_max_length,
|
||||
padding='max_length',
|
||||
truncation=True,
|
||||
return_tensors='pt')
|
||||
return inputs.input_ids
|
||||
|
||||
def compute_time_ids(self, original_size, crops_coords_top_left):
|
||||
target_size = (self.resolution, self.resolution)
|
||||
add_time_ids = list(original_size + crops_coords_top_left
|
||||
+ target_size)
|
||||
add_time_ids = torch.tensor([add_time_ids])
|
||||
add_time_ids = add_time_ids.to(self.device, dtype=self.weight_dtype)
|
||||
return add_time_ids
|
||||
|
||||
def encode_prompt(self,
|
||||
text_encoders,
|
||||
tokenizers,
|
||||
prompt,
|
||||
text_input_ids_list=None):
|
||||
prompt_embeds_list = []
|
||||
|
||||
for i, text_encoder in enumerate(text_encoders):
|
||||
if tokenizers is not None:
|
||||
tokenizer = tokenizers[i]
|
||||
text_input_ids = tokenize_prompt(tokenizer, prompt)
|
||||
else:
|
||||
assert text_input_ids_list is not None
|
||||
text_input_ids = text_input_ids_list[i]
|
||||
|
||||
prompt_embeds = text_encoder(
|
||||
text_input_ids.to(text_encoder.device),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
|
||||
# We are only ALWAYS interested in the pooled output of the final text encoder
|
||||
pooled_prompt_embeds = prompt_embeds[0]
|
||||
prompt_embeds = prompt_embeds.hidden_states[-2]
|
||||
bs_embed, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
|
||||
prompt_embeds_list.append(prompt_embeds)
|
||||
|
||||
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
|
||||
pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
|
||||
return prompt_embeds, pooled_prompt_embeds
|
||||
|
||||
def preprocessing_data(self, text, target):
|
||||
train_crop = transforms.RandomCrop(self.resolution)
|
||||
train_resize = transforms.Resize(
|
||||
self.resolution,
|
||||
interpolation=transforms.InterpolationMode.BILINEAR)
|
||||
train_flip = transforms.RandomHorizontalFlip(p=1.0)
|
||||
image = target
|
||||
original_size = (image.size()[-1], image.size()[-2])
|
||||
image = train_resize(image)
|
||||
y1, x1, h, w = train_crop.get_params(
|
||||
image, (self.resolution, self.resolution))
|
||||
image = crop(image, y1, x1, h, w)
|
||||
if self.random_flip and random.random() < 0.5:
|
||||
# flip
|
||||
x1 = image.size()[-2] - x1
|
||||
image = train_flip(image)
|
||||
crop_top_left = (y1, x1)
|
||||
input_ids_one = self.tokenize_caption(self.tokenizer_one, text)
|
||||
input_ids_two = self.tokenize_caption(self.tokenizer_two, text)
|
||||
|
||||
return original_size, crop_top_left, image, input_ids_one, input_ids_two
|
||||
|
||||
def forward(self, text='', target=None):
|
||||
self.unet.train()
|
||||
self.unet = self.unet.to(self.device)
|
||||
|
||||
# processing data
|
||||
original_size, crop_top_left, image, input_ids_one, input_ids_two = self.preprocessing_data(
|
||||
text, target)
|
||||
# Convert to latent space
|
||||
with torch.no_grad():
|
||||
latents = self.vae.encode(
|
||||
target.to(dtype=self.weight_dtype)).latent_dist.sample()
|
||||
latents = latents * self.vae.config.scaling_factor
|
||||
|
||||
# Sample noise that we'll add to the latents
|
||||
noise = torch.randn_like(latents)
|
||||
bsz = latents.shape[0]
|
||||
# Sample a random timestep for each image
|
||||
timesteps = torch.randint(
|
||||
0,
|
||||
self.noise_scheduler.num_train_timesteps, (bsz, ),
|
||||
device=latents.device)
|
||||
timesteps = timesteps.long()
|
||||
|
||||
# Add noise to the latents according to the noise magnitude at each timestep
|
||||
# (this is the forward diffusion process)
|
||||
noisy_latents = self.noise_scheduler.add_noise(latents, noise,
|
||||
timesteps)
|
||||
|
||||
add_time_ids = self.compute_time_ids(original_size, crop_top_left)
|
||||
|
||||
# Predict the noise residual
|
||||
unet_added_conditions = {'time_ids': add_time_ids}
|
||||
prompt_embeds, pooled_prompt_embeds = self.encode_prompt(
|
||||
text_encoders=[self.text_encoder_one, self.text_encoder_two],
|
||||
tokenizers=None,
|
||||
prompt=None,
|
||||
text_input_ids_list=[input_ids_one, input_ids_two])
|
||||
unet_added_conditions.update({'text_embeds': pooled_prompt_embeds})
|
||||
# Predict the noise residual and compute loss
|
||||
model_pred = self.unet(
|
||||
noisy_latents,
|
||||
timesteps,
|
||||
prompt_embeds,
|
||||
added_cond_kwargs=unet_added_conditions).sample
|
||||
|
||||
# Get the target for loss depending on the prediction type
|
||||
if self.noise_scheduler.config.prediction_type == 'epsilon':
|
||||
target = noise
|
||||
elif self.noise_scheduler.config.prediction_type == 'v_prediction':
|
||||
target = self.noise_scheduler.get_velocity(model_input, noise,
|
||||
timesteps)
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Unknown prediction type {self.noise_scheduler.config.prediction_type}'
|
||||
)
|
||||
|
||||
loss = F.mse_loss(model_pred.float(), target.float(), reduction='mean')
|
||||
|
||||
output = {OutputKeys.LOSS: loss}
|
||||
return output
|
||||
|
||||
def save_pretrained(self,
|
||||
target_folder: Union[str, os.PathLike],
|
||||
save_checkpoint_names: Union[str, List[str]] = None,
|
||||
save_function: Callable = partial(
|
||||
save_checkpoint, with_meta=False),
|
||||
config: Optional[dict] = None,
|
||||
save_config_function: Callable = save_configuration,
|
||||
**kwargs):
|
||||
config['pipeline']['type'] = 'diffusers-stable-diffusion-xl'
|
||||
# Skip copying the original weights for lora and dreambooth method
|
||||
if self.lora_tune or self.dreambooth_tune:
|
||||
pass
|
||||
else:
|
||||
super().save_pretrained(target_folder, save_checkpoint_names,
|
||||
save_function, config,
|
||||
save_config_function, **kwargs)
|
||||
@@ -19,6 +19,9 @@ from modelscope.models.multi_modal.video_synthesis.diffusion import (
|
||||
from modelscope.models.multi_modal.video_synthesis.unet_sd import UNetSD
|
||||
from modelscope.utils.config import Config
|
||||
from modelscope.utils.constant import ModelFile, Tasks
|
||||
from modelscope.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
__all__ = ['TextToVideoSynthesis']
|
||||
|
||||
|
||||
24
modelscope/models/multi_modal/video_to_video/__init__.py
Normal file
24
modelscope/models/multi_modal/video_to_video/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from modelscope.utils.import_utils import LazyImportModule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from .video_to_video_model import VideoToVideo
|
||||
|
||||
else:
|
||||
_import_structure = {
|
||||
'video_to_video_model': ['VideoToVideo'],
|
||||
}
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = LazyImportModule(
|
||||
__name__,
|
||||
globals()['__file__'],
|
||||
_import_structure,
|
||||
module_spec=__spec__,
|
||||
extra_objects={},
|
||||
)
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .autoencoder import *
|
||||
from .embedder import *
|
||||
from .unet_v2v import *
|
||||
@@ -0,0 +1,590 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import collections
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modelscope.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def nonlinearity(x):
|
||||
# swish
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
def Normalize(in_channels, num_groups=32):
|
||||
return torch.nn.GroupNorm(
|
||||
num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def get_first_stage_encoding(encoder_posterior):
|
||||
scale_factor = 0.18215
|
||||
if isinstance(encoder_posterior, DiagonalGaussianDistribution):
|
||||
z = encoder_posterior.sample()
|
||||
elif isinstance(encoder_posterior, torch.Tensor):
|
||||
z = encoder_posterior
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented"
|
||||
)
|
||||
return scale_factor * z
|
||||
|
||||
|
||||
class DiagonalGaussianDistribution(object):
|
||||
|
||||
def __init__(self, parameters, deterministic=False):
|
||||
self.parameters = parameters
|
||||
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
|
||||
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
|
||||
self.deterministic = deterministic
|
||||
self.std = torch.exp(0.5 * self.logvar)
|
||||
self.var = torch.exp(self.logvar)
|
||||
if self.deterministic:
|
||||
self.var = self.std = torch.zeros_like(
|
||||
self.mean).to(device=self.parameters.device)
|
||||
|
||||
def sample(self):
|
||||
x = self.mean + self.std * torch.randn(
|
||||
self.mean.shape).to(device=self.parameters.device)
|
||||
return x
|
||||
|
||||
def kl(self, other=None):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
else:
|
||||
if other is None:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
|
||||
dim=[1, 2, 3])
|
||||
else:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean - other.mean, 2) / other.var
|
||||
+ self.var / other.var - 1.0 - self.logvar + other.logvar,
|
||||
dim=[1, 2, 3])
|
||||
|
||||
def nll(self, sample, dims=[1, 2, 3]):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
logtwopi = np.log(2.0 * np.pi)
|
||||
return 0.5 * torch.sum(
|
||||
logtwopi + self.logvar
|
||||
+ torch.pow(sample - self.mean, 2) / self.var,
|
||||
dim=dims)
|
||||
|
||||
def mode(self):
|
||||
return self.mean
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
in_channels,
|
||||
out_channels=None,
|
||||
conv_shortcut=False,
|
||||
dropout,
|
||||
temb_channels=512):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_conv_shortcut = conv_shortcut
|
||||
|
||||
self.norm1 = Normalize(in_channels)
|
||||
self.conv1 = torch.nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if temb_channels > 0:
|
||||
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
|
||||
self.norm2 = Normalize(out_channels)
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = torch.nn.Conv2d(
|
||||
out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
self.conv_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
else:
|
||||
self.nin_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0)
|
||||
|
||||
def forward(self, x, temb):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
if temb is not None:
|
||||
h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
|
||||
|
||||
h = self.norm2(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.dropout(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
x = self.conv_shortcut(x)
|
||||
else:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = Normalize(in_channels)
|
||||
self.q = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.k = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.v = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.proj_out = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h * w)
|
||||
q = q.permute(0, 2, 1)
|
||||
k = k.reshape(b, c, h * w)
|
||||
w_ = torch.bmm(q, k)
|
||||
w_ = w_ * (int(c)**(-0.5))
|
||||
w_ = torch.nn.functional.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h * w)
|
||||
w_ = w_.permute(0, 2, 1)
|
||||
h_ = torch.bmm(v, w_)
|
||||
h_ = h_.reshape(b, c, h, w)
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x + h_
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.nn.functional.interpolate(
|
||||
x, scale_factor=2.0, mode='nearest')
|
||||
if self.with_conv:
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
if self.with_conv:
|
||||
pad = (0, 1, 0, 1)
|
||||
x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
|
||||
x = self.conv(x)
|
||||
else:
|
||||
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
double_z=True,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignore_kwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
|
||||
# downsampling
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1, ) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in,
|
||||
2 * z_channels if double_z else z_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1], temb)
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
give_pre_end=False,
|
||||
tanh_out=False,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignorekwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.give_pre_end = give_pre_end
|
||||
self.tanh_out = tanh_out
|
||||
|
||||
# compute block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2**(self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
logger.info('Working with z of shape {} = {} dimensions.'.format(
|
||||
self.z_shape, np.prod(self.z_shape)))
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up)
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z):
|
||||
self.last_z_shape = z.shape
|
||||
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h, temb)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
if self.give_pre_end:
|
||||
return h
|
||||
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
if self.tanh_out:
|
||||
h = torch.tanh(h)
|
||||
return h
|
||||
|
||||
|
||||
class AutoencoderKL(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
embed_dim,
|
||||
pretrained=None,
|
||||
ignore_keys=[],
|
||||
image_key='image',
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
ema_decay=None,
|
||||
learn_logvar=False,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
self.learn_logvar = learn_logvar
|
||||
self.image_key = image_key
|
||||
self.encoder = Encoder(**ddconfig)
|
||||
self.decoder = Decoder(**ddconfig)
|
||||
assert ddconfig['double_z']
|
||||
self.quant_conv = torch.nn.Conv2d(2 * ddconfig['z_channels'],
|
||||
2 * embed_dim, 1)
|
||||
self.post_quant_conv = torch.nn.Conv2d(embed_dim,
|
||||
ddconfig['z_channels'], 1)
|
||||
self.embed_dim = embed_dim
|
||||
if colorize_nlabels is not None:
|
||||
assert type(colorize_nlabels) == int
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, colorize_nlabels, 1, 1))
|
||||
if monitor is not None:
|
||||
self.monitor = monitor
|
||||
|
||||
self.use_ema = ema_decay is not None
|
||||
|
||||
if pretrained is not None:
|
||||
self.init_from_ckpt(pretrained, ignore_keys=ignore_keys)
|
||||
|
||||
def init_from_ckpt(self, path, ignore_keys=list()):
|
||||
sd = torch.load(path, map_location='cpu')['state_dict']
|
||||
keys = list(sd.keys())
|
||||
sd_new = collections.OrderedDict()
|
||||
for k in keys:
|
||||
if k.find('first_stage_model') >= 0:
|
||||
k_new = k.split('first_stage_model.')[-1]
|
||||
sd_new[k_new] = sd[k]
|
||||
self.load_state_dict(sd_new, strict=True)
|
||||
logger.info(f'Restored from {path}')
|
||||
|
||||
def on_train_batch_end(self, *args, **kwargs):
|
||||
if self.use_ema:
|
||||
self.model_ema(self)
|
||||
|
||||
def encode(self, x):
|
||||
h = self.encoder(x)
|
||||
moments = self.quant_conv(h)
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
return posterior
|
||||
|
||||
def decode(self, z):
|
||||
z = self.post_quant_conv(z)
|
||||
dec = self.decoder(z)
|
||||
return dec
|
||||
|
||||
def forward(self, input, sample_posterior=True):
|
||||
posterior = self.encode(input)
|
||||
if sample_posterior:
|
||||
z = posterior.sample()
|
||||
else:
|
||||
z = posterior.mode()
|
||||
dec = self.decode(z)
|
||||
return dec, posterior
|
||||
|
||||
def get_input(self, batch, k):
|
||||
x = batch[k]
|
||||
if len(x.shape) == 3:
|
||||
x = x[..., None]
|
||||
x = x.permute(0, 3, 1,
|
||||
2).to(memory_format=torch.contiguous_format).float()
|
||||
return x
|
||||
|
||||
def get_last_layer(self):
|
||||
return self.decoder.conv_out.weight
|
||||
|
||||
@torch.no_grad()
|
||||
def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
if not only_inputs:
|
||||
xrec, posterior = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec.shape[1] > 3
|
||||
x = self.to_rgb(x)
|
||||
xrec = self.to_rgb(xrec)
|
||||
log['samples'] = self.decode(torch.randn_like(posterior.sample()))
|
||||
log['reconstructions'] = xrec
|
||||
if log_ema or self.use_ema:
|
||||
with self.ema_scope():
|
||||
xrec_ema, posterior_ema = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec_ema.shape[1] > 3
|
||||
xrec_ema = self.to_rgb(xrec_ema)
|
||||
log['samples_ema'] = self.decode(
|
||||
torch.randn_like(posterior_ema.sample()))
|
||||
log['reconstructions_ema'] = xrec_ema
|
||||
log['inputs'] = x
|
||||
return log
|
||||
|
||||
def to_rgb(self, x):
|
||||
assert self.image_key == 'segmentation'
|
||||
if not hasattr(self, 'colorize'):
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, x.shape[1], 1, 1).to(x))
|
||||
x = F.conv2d(x, weight=self.colorize)
|
||||
x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
|
||||
return x
|
||||
|
||||
|
||||
class IdentityFirstStage(torch.nn.Module):
|
||||
|
||||
def __init__(self, *args, vq_interface=False, **kwargs):
|
||||
self.vq_interface = vq_interface
|
||||
super().__init__()
|
||||
|
||||
def encode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def decode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def quantize(self, x, *args, **kwargs):
|
||||
if self.vq_interface:
|
||||
return x, None, [None, None, None]
|
||||
return x
|
||||
|
||||
def forward(self, x, *args, **kwargs):
|
||||
return x
|
||||
@@ -0,0 +1,76 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as T
|
||||
|
||||
|
||||
class FrozenOpenCLIPEmbedder(nn.Module):
|
||||
"""
|
||||
Uses the OpenCLIP transformer encoder for text
|
||||
"""
|
||||
LAYERS = ['last', 'penultimate']
|
||||
|
||||
def __init__(self,
|
||||
pretrained,
|
||||
arch='ViT-H-14',
|
||||
device='cuda',
|
||||
max_length=77,
|
||||
freeze=True,
|
||||
layer='penultimate'):
|
||||
super().__init__()
|
||||
assert layer in self.LAYERS
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=pretrained)
|
||||
|
||||
del model.visual
|
||||
self.model = model
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
|
||||
if freeze:
|
||||
self.freeze()
|
||||
self.layer = layer
|
||||
if self.layer == 'last':
|
||||
self.layer_idx = 0
|
||||
elif self.layer == 'penultimate':
|
||||
self.layer_idx = 1
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
def freeze(self):
|
||||
self.model = self.model.eval()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, text):
|
||||
tokens = open_clip.tokenize(text)
|
||||
z = self.encode_with_transformer(tokens.to(self.device))
|
||||
return z
|
||||
|
||||
def encode_with_transformer(self, text):
|
||||
x = self.model.token_embedding(text)
|
||||
x = x + self.model.positional_embedding
|
||||
x = x.permute(1, 0, 2)
|
||||
x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
|
||||
x = x.permute(1, 0, 2)
|
||||
x = self.model.ln_final(x)
|
||||
return x
|
||||
|
||||
def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
|
||||
for i, r in enumerate(self.model.transformer.resblocks):
|
||||
if i == len(self.model.transformer.resblocks) - self.layer_idx:
|
||||
break
|
||||
if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(
|
||||
):
|
||||
x = checkpoint(r, x, attn_mask)
|
||||
else:
|
||||
x = r(x, attn_mask=attn_mask)
|
||||
return x
|
||||
|
||||
def encode(self, text):
|
||||
return self(text)
|
||||
1542
modelscope/models/multi_modal/video_to_video/modules/unet_v2v.py
Normal file
1542
modelscope/models/multi_modal/video_to_video/modules/unet_v2v.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import os
|
||||
171
modelscope/models/multi_modal/video_to_video/utils/config.py
Normal file
171
modelscope/models/multi_modal/video_to_video/utils/config.py
Normal file
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import os.path as osp
|
||||
from datetime import datetime
|
||||
|
||||
import torch
|
||||
from easydict import EasyDict
|
||||
|
||||
cfg = EasyDict(__name__='Config: VideoLDM Decoder')
|
||||
|
||||
# ---------------------------work dir--------------------------
|
||||
cfg.work_dir = 'workspace/'
|
||||
|
||||
# ---------------------------Global Variable-----------------------------------
|
||||
cfg.resolution = [448, 256]
|
||||
cfg.max_frames = 32
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Dataset Parameter---------------------------------
|
||||
cfg.mean = [0.5, 0.5, 0.5]
|
||||
cfg.std = [0.5, 0.5, 0.5]
|
||||
cfg.max_words = 1000
|
||||
|
||||
# PlaceHolder
|
||||
cfg.vit_out_dim = 1024
|
||||
cfg.vit_resolution = [224, 224]
|
||||
cfg.depth_clamp = 10.0
|
||||
cfg.misc_size = 384
|
||||
cfg.depth_std = 20.0
|
||||
|
||||
cfg.frame_lens = 32
|
||||
cfg.sample_fps = 8
|
||||
|
||||
cfg.batch_sizes = 1
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Mode Parameters-----------------------------------
|
||||
# Diffusion
|
||||
cfg.schedule = 'cosine'
|
||||
cfg.num_timesteps = 1000
|
||||
cfg.mean_type = 'v'
|
||||
cfg.var_type = 'fixed_small'
|
||||
cfg.loss_type = 'mse'
|
||||
cfg.ddim_timesteps = 50
|
||||
cfg.ddim_eta = 0.0
|
||||
cfg.clamp = 1.0
|
||||
cfg.share_noise = False
|
||||
cfg.use_div_loss = False
|
||||
cfg.noise_strength = 0.1
|
||||
|
||||
# classifier-free guidance
|
||||
cfg.p_zero = 0.1
|
||||
cfg.guide_scale = 3.0
|
||||
|
||||
# clip vision encoder
|
||||
cfg.vit_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||
cfg.vit_std = [0.26862954, 0.26130258, 0.27577711]
|
||||
|
||||
# Model
|
||||
cfg.scale_factor = 0.18215
|
||||
cfg.use_fp16 = True
|
||||
cfg.temporal_attention = True
|
||||
cfg.decoder_bs = 8
|
||||
|
||||
cfg.UNet = {
|
||||
'type': 'Vid2VidSDUNet',
|
||||
'in_dim': 4,
|
||||
'dim': 320,
|
||||
'y_dim': cfg.vit_out_dim,
|
||||
'context_dim': 1024,
|
||||
'out_dim': 8 if cfg.var_type.startswith('learned') else 4,
|
||||
'dim_mult': [1, 2, 4, 4],
|
||||
'num_heads': 8,
|
||||
'head_dim': 64,
|
||||
'num_res_blocks': 2,
|
||||
'attn_scales': [1 / 1, 1 / 2, 1 / 4],
|
||||
'dropout': 0.1,
|
||||
'temporal_attention': cfg.temporal_attention,
|
||||
'temporal_attn_times': 1,
|
||||
'use_checkpoint': False,
|
||||
'use_fps_condition': False,
|
||||
'use_sim_mask': False,
|
||||
'num_tokens': 4,
|
||||
'default_fps': 8,
|
||||
'input_dim': 1024
|
||||
}
|
||||
|
||||
cfg.guidances = []
|
||||
|
||||
# auotoencoder from stabel diffusion
|
||||
cfg.auto_encoder = {
|
||||
'type': 'AutoencoderKL',
|
||||
'ddconfig': {
|
||||
'double_z': True,
|
||||
'z_channels': 4,
|
||||
'resolution': 256,
|
||||
'in_channels': 3,
|
||||
'out_ch': 3,
|
||||
'ch': 128,
|
||||
'ch_mult': [1, 2, 4, 4],
|
||||
'num_res_blocks': 2,
|
||||
'attn_resolutions': [],
|
||||
'dropout': 0.0
|
||||
},
|
||||
'embed_dim': 4,
|
||||
'pretrained': 'models/v2-1_512-ema-pruned.ckpt'
|
||||
}
|
||||
# clip embedder
|
||||
cfg.embedder = {
|
||||
'type': 'FrozenOpenCLIPEmbedder',
|
||||
'layer': 'penultimate',
|
||||
'vit_resolution': [224, 224],
|
||||
'pretrained': 'open_clip_pytorch_model.bin'
|
||||
}
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Training Settings---------------------------------
|
||||
# training and optimizer
|
||||
cfg.ema_decay = 0.9999
|
||||
cfg.num_steps = 600000
|
||||
cfg.lr = 5e-5
|
||||
cfg.weight_decay = 0.0
|
||||
cfg.betas = (0.9, 0.999)
|
||||
cfg.eps = 1.0e-8
|
||||
cfg.chunk_size = 16
|
||||
cfg.alpha = 0.7
|
||||
cfg.save_ckp_interval = 1000
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ----------------------------Pretrain Settings---------------------------------
|
||||
# Default: load 2d pretrain
|
||||
cfg.fix_weight = False
|
||||
cfg.load_match = False
|
||||
cfg.pretrained_checkpoint = 'v2-1_512-ema-pruned.ckpt'
|
||||
cfg.pretrained_image_keys = 'stable_diffusion_image_key_temporal_attention_x1.json'
|
||||
cfg.resume_checkpoint = 'img2video_ldm_0779000.pth'
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# -----------------------------Visual-------------------------------------------
|
||||
# Visual videos
|
||||
cfg.viz_interval = 1000
|
||||
cfg.visual_train = {
|
||||
'type': 'VisualVideoTextDuringTrain',
|
||||
}
|
||||
cfg.visual_inference = {
|
||||
'type': 'VisualGeneratedVideos',
|
||||
}
|
||||
cfg.inference_list_path = ''
|
||||
|
||||
# logging
|
||||
cfg.log_interval = 100
|
||||
|
||||
# Default log_dir
|
||||
cfg.log_dir = 'workspace/output_data'
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------Others--------------------------------------------
|
||||
# seed
|
||||
cfg.seed = 8888
|
||||
cfg.negative_prompt = 'worst quality, normal quality, low quality, low res, blurry, text, \
|
||||
watermark, logo, banner, extra digits, cropped, jpeg artifacts, signature, username, error, \
|
||||
sketch ,duplicate, ugly, monochrome, horror, geometry, mutation, disgusting'
|
||||
|
||||
cfg.positive_prompt = ', cinematic, High Contrast, highly detailed, unreal engine, \
|
||||
taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, \
|
||||
32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, \
|
||||
hyper sharpness, perfect without deformations, Unreal Engine 5, 4k render'
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,247 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import random
|
||||
|
||||
import torch
|
||||
|
||||
from .schedules_sdedit import karras_schedule
|
||||
from .solvers_sdedit import sample_dpmpp_2m_sde, sample_heun
|
||||
|
||||
__all__ = ['GaussianDiffusion_SDEdit']
|
||||
|
||||
|
||||
def _i(tensor, t, x):
|
||||
shape = (x.size(0), ) + (1, ) * (x.ndim - 1)
|
||||
return tensor[t.to(tensor.device)].view(shape).to(x.device)
|
||||
|
||||
|
||||
class GaussianDiffusion_SDEdit(object):
|
||||
|
||||
def __init__(self, sigmas, prediction_type='eps'):
|
||||
assert prediction_type in {'x0', 'eps', 'v'}
|
||||
self.sigmas = sigmas
|
||||
self.alphas = torch.sqrt(1 - sigmas**2)
|
||||
self.num_timesteps = len(sigmas)
|
||||
self.prediction_type = prediction_type
|
||||
|
||||
def diffuse(self, x0, t, noise=None):
|
||||
noise = torch.randn_like(x0) if noise is None else noise
|
||||
xt = _i(self.alphas, t, x0) * x0 + _i(self.sigmas, t, x0) * noise
|
||||
return xt
|
||||
|
||||
def denoise(self,
|
||||
xt,
|
||||
t,
|
||||
s,
|
||||
model,
|
||||
model_kwargs={},
|
||||
guide_scale=None,
|
||||
guide_rescale=None,
|
||||
clamp=None,
|
||||
percentile=None):
|
||||
s = t - 1 if s is None else s
|
||||
|
||||
# hyperparams
|
||||
sigmas = _i(self.sigmas, t, xt)
|
||||
alphas = _i(self.alphas, t, xt)
|
||||
alphas_s = _i(self.alphas, s.clamp(0), xt)
|
||||
alphas_s[s < 0] = 1.
|
||||
sigmas_s = torch.sqrt(1 - alphas_s**2)
|
||||
|
||||
# precompute variables
|
||||
betas = 1 - (alphas / alphas_s)**2
|
||||
coef1 = betas * alphas_s / sigmas**2
|
||||
coef2 = (alphas * sigmas_s**2) / (alphas_s * sigmas**2)
|
||||
var = betas * (sigmas_s / sigmas)**2
|
||||
log_var = torch.log(var).clamp_(-20, 20)
|
||||
|
||||
# prediction
|
||||
if guide_scale is None:
|
||||
assert isinstance(model_kwargs, dict)
|
||||
out = model(xt, t=t, **model_kwargs)
|
||||
else:
|
||||
# classifier-free guidance
|
||||
assert isinstance(model_kwargs, list) and len(model_kwargs) == 2
|
||||
y_out = model(xt, t=t, **model_kwargs[0])
|
||||
if guide_scale == 1.:
|
||||
out = y_out
|
||||
else:
|
||||
u_out = model(xt, t=t, **model_kwargs[1])
|
||||
out = u_out + guide_scale * (y_out - u_out)
|
||||
|
||||
if guide_rescale is not None:
|
||||
assert guide_rescale >= 0 and guide_rescale <= 1
|
||||
ratio = (
|
||||
y_out.flatten(1).std(dim=1) / # noqa
|
||||
(out.flatten(1).std(dim=1) + 1e-12)
|
||||
).view((-1, ) + (1, ) * (y_out.ndim - 1))
|
||||
out *= guide_rescale * ratio + (1 - guide_rescale) * 1.0
|
||||
|
||||
# compute x0
|
||||
if self.prediction_type == 'x0':
|
||||
x0 = out
|
||||
elif self.prediction_type == 'eps':
|
||||
x0 = (xt - sigmas * out) / alphas
|
||||
elif self.prediction_type == 'v':
|
||||
x0 = alphas * xt - sigmas * out
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'prediction_type {self.prediction_type} not implemented')
|
||||
|
||||
# restrict the range of x0
|
||||
if percentile is not None:
|
||||
assert percentile > 0 and percentile <= 1
|
||||
s = torch.quantile(x0.flatten(1).abs(), percentile, dim=1)
|
||||
s = s.clamp_(1.0).view((-1, ) + (1, ) * (xt.ndim - 1))
|
||||
x0 = torch.min(s, torch.max(-s, x0)) / s
|
||||
elif clamp is not None:
|
||||
x0 = x0.clamp(-clamp, clamp)
|
||||
|
||||
# recompute eps using the restricted x0
|
||||
eps = (xt - alphas * x0) / sigmas
|
||||
|
||||
# compute mu (mean of posterior distribution) using the restricted x0
|
||||
mu = coef1 * x0 + coef2 * xt
|
||||
return mu, var, log_var, x0, eps
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self,
|
||||
noise,
|
||||
model,
|
||||
model_kwargs={},
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
guide_rescale=None,
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
solver='euler_a',
|
||||
steps=20,
|
||||
t_max=None,
|
||||
t_min=None,
|
||||
discretization=None,
|
||||
discard_penultimate_step=None,
|
||||
return_intermediate=None,
|
||||
show_progress=False,
|
||||
seed=-1,
|
||||
**kwargs):
|
||||
# sanity check
|
||||
assert isinstance(steps, (int, torch.LongTensor))
|
||||
assert t_max is None or (t_max > 0 and t_max <= self.num_timesteps - 1)
|
||||
assert t_min is None or (t_min >= 0 and t_min < self.num_timesteps - 1)
|
||||
assert discretization in (None, 'leading', 'linspace', 'trailing')
|
||||
assert discard_penultimate_step in (None, True, False)
|
||||
assert return_intermediate in (None, 'x0', 'xt')
|
||||
|
||||
# function of diffusion solver
|
||||
solver_fn = {
|
||||
'heun': sample_heun,
|
||||
'dpmpp_2m_sde': sample_dpmpp_2m_sde
|
||||
}[solver]
|
||||
|
||||
# options
|
||||
schedule = 'karras' if 'karras' in solver else None
|
||||
discretization = discretization or 'linspace'
|
||||
seed = seed if seed >= 0 else random.randint(0, 2**31)
|
||||
if isinstance(steps, torch.LongTensor):
|
||||
discard_penultimate_step = False
|
||||
if discard_penultimate_step is None:
|
||||
discard_penultimate_step = True if solver in (
|
||||
'dpm2', 'dpm2_ancestral', 'dpmpp_2m_sde', 'dpm2_karras',
|
||||
'dpm2_ancestral_karras', 'dpmpp_2m_sde_karras') else False
|
||||
|
||||
# function for denoising xt to get x0
|
||||
intermediates = []
|
||||
|
||||
def model_fn(xt, sigma):
|
||||
# denoising
|
||||
t = self._sigma_to_t(sigma).repeat(len(xt)).round().long()
|
||||
x0 = self.denoise(xt, t, None, model, model_kwargs, guide_scale,
|
||||
guide_rescale, clamp, percentile)[-2]
|
||||
|
||||
# collect intermediate outputs
|
||||
if return_intermediate == 'xt':
|
||||
intermediates.append(xt)
|
||||
elif return_intermediate == 'x0':
|
||||
intermediates.append(x0)
|
||||
return x0
|
||||
|
||||
# get timesteps
|
||||
if isinstance(steps, int):
|
||||
steps += 1 if discard_penultimate_step else 0
|
||||
t_max = self.num_timesteps - 1 if t_max is None else t_max
|
||||
t_min = 0 if t_min is None else t_min
|
||||
|
||||
# discretize timesteps
|
||||
if discretization == 'leading':
|
||||
steps = torch.arange(t_min, t_max + 1,
|
||||
(t_max - t_min + 1) / steps).flip(0)
|
||||
elif discretization == 'linspace':
|
||||
steps = torch.linspace(t_max, t_min, steps)
|
||||
elif discretization == 'trailing':
|
||||
steps = torch.arange(t_max, t_min - 1,
|
||||
-((t_max - t_min + 1) / steps))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'{discretization} discretization not implemented')
|
||||
steps = steps.clamp_(t_min, t_max)
|
||||
steps = torch.as_tensor(
|
||||
steps, dtype=torch.float32, device=noise.device)
|
||||
|
||||
# get sigmas
|
||||
sigmas = self._t_to_sigma(steps)
|
||||
sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
|
||||
if schedule == 'karras':
|
||||
if sigmas[0] == float('inf'):
|
||||
sigmas = karras_schedule(
|
||||
n=len(steps) - 1,
|
||||
sigma_min=sigmas[sigmas > 0].min().item(),
|
||||
sigma_max=sigmas[sigmas < float('inf')].max().item(),
|
||||
rho=7.).to(sigmas)
|
||||
sigmas = torch.cat([
|
||||
sigmas.new_tensor([float('inf')]), sigmas,
|
||||
sigmas.new_zeros([1])
|
||||
])
|
||||
else:
|
||||
sigmas = karras_schedule(
|
||||
n=len(steps),
|
||||
sigma_min=sigmas[sigmas > 0].min().item(),
|
||||
sigma_max=sigmas.max().item(),
|
||||
rho=7.).to(sigmas)
|
||||
sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
|
||||
if discard_penultimate_step:
|
||||
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
|
||||
|
||||
# sampling
|
||||
x0 = solver_fn(
|
||||
noise, model_fn, sigmas, show_progress=show_progress, **kwargs)
|
||||
return (x0, intermediates) if return_intermediate is not None else x0
|
||||
|
||||
def _sigma_to_t(self, sigma):
|
||||
if sigma == float('inf'):
|
||||
t = torch.full_like(sigma, len(self.sigmas) - 1)
|
||||
else:
|
||||
log_sigmas = torch.sqrt(self.sigmas**2 / # noqa
|
||||
(1 - self.sigmas**2)).log().to(sigma)
|
||||
log_sigma = sigma.log()
|
||||
dists = log_sigma - log_sigmas[:, None]
|
||||
low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(
|
||||
max=log_sigmas.shape[0] - 2)
|
||||
high_idx = low_idx + 1
|
||||
low, high = log_sigmas[low_idx], log_sigmas[high_idx]
|
||||
w = (low - log_sigma) / (low - high)
|
||||
w = w.clamp(0, 1)
|
||||
t = (1 - w) * low_idx + w * high_idx
|
||||
t = t.view(sigma.shape)
|
||||
if t.ndim == 0:
|
||||
t = t.unsqueeze(0)
|
||||
return t
|
||||
|
||||
def _t_to_sigma(self, t):
|
||||
t = t.float()
|
||||
low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
|
||||
log_sigmas = torch.sqrt(self.sigmas**2 / # noqa
|
||||
(1 - self.sigmas**2)).log().to(t)
|
||||
log_sigma = (1 - w) * log_sigmas[low_idx] + w * log_sigmas[high_idx]
|
||||
log_sigma[torch.isnan(log_sigma)
|
||||
| torch.isinf(log_sigma)] = float('inf')
|
||||
return log_sigma.exp()
|
||||
@@ -0,0 +1,85 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def betas_to_sigmas(betas):
|
||||
return torch.sqrt(1 - torch.cumprod(1 - betas, dim=0))
|
||||
|
||||
|
||||
def sigmas_to_betas(sigmas):
|
||||
square_alphas = 1 - sigmas**2
|
||||
betas = 1 - torch.cat(
|
||||
[square_alphas[:1], square_alphas[1:] / square_alphas[:-1]])
|
||||
return betas
|
||||
|
||||
|
||||
def logsnrs_to_sigmas(logsnrs):
|
||||
return torch.sqrt(torch.sigmoid(-logsnrs))
|
||||
|
||||
|
||||
def sigmas_to_logsnrs(sigmas):
|
||||
square_sigmas = sigmas**2
|
||||
return torch.log(square_sigmas / (1 - square_sigmas))
|
||||
|
||||
|
||||
def _logsnr_cosine(n, logsnr_min=-15, logsnr_max=15):
|
||||
t_min = math.atan(math.exp(-0.5 * logsnr_min))
|
||||
t_max = math.atan(math.exp(-0.5 * logsnr_max))
|
||||
t = torch.linspace(1, 0, n)
|
||||
logsnrs = -2 * torch.log(torch.tan(t_min + t * (t_max - t_min)))
|
||||
return logsnrs
|
||||
|
||||
|
||||
def _logsnr_cosine_shifted(n, logsnr_min=-15, logsnr_max=15, scale=2):
|
||||
logsnrs = _logsnr_cosine(n, logsnr_min, logsnr_max)
|
||||
logsnrs += 2 * math.log(1 / scale)
|
||||
return logsnrs
|
||||
|
||||
|
||||
def _logsnr_cosine_interp(n,
|
||||
logsnr_min=-15,
|
||||
logsnr_max=15,
|
||||
scale_min=2,
|
||||
scale_max=4):
|
||||
t = torch.linspace(1, 0, n)
|
||||
logsnrs_min = _logsnr_cosine_shifted(n, logsnr_min, logsnr_max, scale_min)
|
||||
logsnrs_max = _logsnr_cosine_shifted(n, logsnr_min, logsnr_max, scale_max)
|
||||
logsnrs = t * logsnrs_min + (1 - t) * logsnrs_max
|
||||
return logsnrs
|
||||
|
||||
|
||||
def karras_schedule(n, sigma_min=0.002, sigma_max=80.0, rho=7.0):
|
||||
ramp = torch.linspace(1, 0, n)
|
||||
min_inv_rho = sigma_min**(1 / rho)
|
||||
max_inv_rho = sigma_max**(1 / rho)
|
||||
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
|
||||
sigmas = torch.sqrt(sigmas**2 / (1 + sigmas**2))
|
||||
return sigmas
|
||||
|
||||
|
||||
def logsnr_cosine_interp_schedule(n,
|
||||
logsnr_min=-15,
|
||||
logsnr_max=15,
|
||||
scale_min=2,
|
||||
scale_max=4):
|
||||
return logsnrs_to_sigmas(
|
||||
_logsnr_cosine_interp(n, logsnr_min, logsnr_max, scale_min, scale_max))
|
||||
|
||||
|
||||
def noise_schedule(schedule='logsnr_cosine_interp',
|
||||
n=1000,
|
||||
zero_terminal_snr=False,
|
||||
**kwargs):
|
||||
# compute sigmas
|
||||
sigmas = {
|
||||
'logsnr_cosine_interp': logsnr_cosine_interp_schedule
|
||||
}[schedule](n, **kwargs)
|
||||
|
||||
# post-processing
|
||||
if zero_terminal_snr and sigmas.max() != 1.0:
|
||||
scale = (1.0 - sigmas.min()) / (sigmas.max() - sigmas.min())
|
||||
sigmas = sigmas.min() + scale * (sigmas - sigmas.min())
|
||||
return sigmas
|
||||
14
modelscope/models/multi_modal/video_to_video/utils/seed.py
Normal file
14
modelscope/models/multi_modal/video_to_video/utils/seed.py
Normal file
@@ -0,0 +1,14 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def setup_seed(seed):
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
@@ -0,0 +1,194 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import torch
|
||||
import torchsde
|
||||
from tqdm.auto import trange
|
||||
|
||||
|
||||
def get_ancestral_step(sigma_from, sigma_to, eta=1.):
|
||||
"""
|
||||
Calculates the noise level (sigma_down) to step down to and the amount
|
||||
of noise to add (sigma_up) when doing an ancestral sampling step.
|
||||
"""
|
||||
if not eta:
|
||||
return sigma_to, 0.
|
||||
sigma_up = min(
|
||||
sigma_to,
|
||||
eta * (
|
||||
sigma_to**2 * # noqa
|
||||
(sigma_from**2 - sigma_to**2) / sigma_from**2)**0.5)
|
||||
sigma_down = (sigma_to**2 - sigma_up**2)**0.5
|
||||
return sigma_down, sigma_up
|
||||
|
||||
|
||||
def get_scalings(sigma):
|
||||
c_out = -sigma
|
||||
c_in = 1 / (sigma**2 + 1.**2)**0.5
|
||||
return c_out, c_in
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_heun(noise,
|
||||
model,
|
||||
sigmas,
|
||||
s_churn=0.,
|
||||
s_tmin=0.,
|
||||
s_tmax=float('inf'),
|
||||
s_noise=1.,
|
||||
show_progress=True):
|
||||
"""
|
||||
Implements Algorithm 2 (Heun steps) from Karras et al. (2022).
|
||||
"""
|
||||
x = noise * sigmas[0]
|
||||
for i in trange(len(sigmas) - 1, disable=not show_progress):
|
||||
gamma = 0.
|
||||
if s_tmin <= sigmas[i] <= s_tmax and sigmas[i] < float('inf'):
|
||||
gamma = min(s_churn / (len(sigmas) - 1), 2**0.5 - 1)
|
||||
eps = torch.randn_like(x) * s_noise
|
||||
sigma_hat = sigmas[i] * (gamma + 1)
|
||||
if gamma > 0:
|
||||
x = x + eps * (sigma_hat**2 - sigmas[i]**2)**0.5
|
||||
if sigmas[i] == float('inf'):
|
||||
# Euler method
|
||||
denoised = model(noise, sigma_hat)
|
||||
x = denoised + sigmas[i + 1] * (gamma + 1) * noise
|
||||
else:
|
||||
_, c_in = get_scalings(sigma_hat)
|
||||
denoised = model(x * c_in, sigma_hat)
|
||||
d = (x - denoised) / sigma_hat
|
||||
dt = sigmas[i + 1] - sigma_hat
|
||||
if sigmas[i + 1] == 0:
|
||||
# Euler method
|
||||
x = x + d * dt
|
||||
else:
|
||||
# Heun's method
|
||||
x_2 = x + d * dt
|
||||
_, c_in = get_scalings(sigmas[i + 1])
|
||||
denoised_2 = model(x_2 * c_in, sigmas[i + 1])
|
||||
d_2 = (x_2 - denoised_2) / sigmas[i + 1]
|
||||
d_prime = (d + d_2) / 2
|
||||
x = x + d_prime * dt
|
||||
return x
|
||||
|
||||
|
||||
class BatchedBrownianTree:
|
||||
"""
|
||||
A wrapper around torchsde.BrownianTree that enables batches of entropy.
|
||||
"""
|
||||
|
||||
def __init__(self, x, t0, t1, seed=None, **kwargs):
|
||||
t0, t1, self.sign = self.sort(t0, t1)
|
||||
w0 = kwargs.get('w0', torch.zeros_like(x))
|
||||
if seed is None:
|
||||
seed = torch.randint(0, 2**63 - 1, []).item()
|
||||
self.batched = True
|
||||
try:
|
||||
assert len(seed) == x.shape[0]
|
||||
w0 = w0[0]
|
||||
except TypeError:
|
||||
seed = [seed]
|
||||
self.batched = False
|
||||
self.trees = [
|
||||
torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs)
|
||||
for s in seed
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def sort(a, b):
|
||||
return (a, b, 1) if a < b else (b, a, -1)
|
||||
|
||||
def __call__(self, t0, t1):
|
||||
t0, t1, sign = self.sort(t0, t1)
|
||||
w = torch.stack([tree(t0, t1) for tree in self.trees]) * (
|
||||
self.sign * sign)
|
||||
return w if self.batched else w[0]
|
||||
|
||||
|
||||
class BrownianTreeNoiseSampler:
|
||||
"""
|
||||
A noise sampler backed by a torchsde.BrownianTree.
|
||||
|
||||
Args:
|
||||
x (Tensor): The tensor whose shape, device and dtype to use to generate
|
||||
random samples.
|
||||
sigma_min (float): The low end of the valid interval.
|
||||
sigma_max (float): The high end of the valid interval.
|
||||
seed (int or List[int]): The random seed. If a list of seeds is
|
||||
supplied instead of a single integer, then the noise sampler will
|
||||
use one BrownianTree per batch item, each with its own seed.
|
||||
transform (callable): A function that maps sigma to the sampler's
|
||||
internal timestep.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
x,
|
||||
sigma_min,
|
||||
sigma_max,
|
||||
seed=None,
|
||||
transform=lambda x: x):
|
||||
self.transform = transform
|
||||
t0 = self.transform(torch.as_tensor(sigma_min))
|
||||
t1 = self.transform(torch.as_tensor(sigma_max))
|
||||
self.tree = BatchedBrownianTree(x, t0, t1, seed)
|
||||
|
||||
def __call__(self, sigma, sigma_next):
|
||||
t0 = self.transform(torch.as_tensor(sigma))
|
||||
t1 = self.transform(torch.as_tensor(sigma_next))
|
||||
return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_dpmpp_2m_sde(noise,
|
||||
model,
|
||||
sigmas,
|
||||
eta=1.,
|
||||
s_noise=1.,
|
||||
solver_type='midpoint',
|
||||
show_progress=True):
|
||||
"""
|
||||
DPM-Solver++ (2M) SDE.
|
||||
"""
|
||||
assert solver_type in {'heun', 'midpoint'}
|
||||
|
||||
x = noise * sigmas[0]
|
||||
sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas[
|
||||
sigmas < float('inf')].max()
|
||||
noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max)
|
||||
old_denoised = None
|
||||
h_last = None
|
||||
|
||||
for i in trange(len(sigmas) - 1, disable=not show_progress):
|
||||
if sigmas[i] == float('inf'):
|
||||
# Euler method
|
||||
denoised = model(noise, sigmas[i])
|
||||
x = denoised + sigmas[i + 1] * noise
|
||||
else:
|
||||
_, c_in = get_scalings(sigmas[i])
|
||||
denoised = model(x * c_in, sigmas[i])
|
||||
if sigmas[i + 1] == 0:
|
||||
# Denoising step
|
||||
x = denoised
|
||||
else:
|
||||
# DPM-Solver++(2M) SDE
|
||||
t, s = -sigmas[i].log(), -sigmas[i + 1].log()
|
||||
h = s - t
|
||||
eta_h = eta * h
|
||||
|
||||
x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + \
|
||||
(-h - eta_h).expm1().neg() * denoised
|
||||
|
||||
if old_denoised is not None:
|
||||
r = h_last / h
|
||||
if solver_type == 'heun':
|
||||
x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * \
|
||||
(1 / r) * (denoised - old_denoised)
|
||||
elif solver_type == 'midpoint':
|
||||
x = x + 0.5 * (-h - eta_h).expm1().neg() * \
|
||||
(1 / r) * (denoised - old_denoised)
|
||||
|
||||
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[
|
||||
i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
|
||||
|
||||
old_denoised = denoised
|
||||
h_last = h
|
||||
return x
|
||||
404
modelscope/models/multi_modal/video_to_video/utils/transforms.py
Normal file
404
modelscope/models/multi_modal/video_to_video/utils/transforms.py
Normal file
@@ -0,0 +1,404 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as F
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
__all__ = [
|
||||
'Compose', 'Resize', 'Rescale', 'CenterCrop', 'CenterCropV2',
|
||||
'CenterCropWide', 'RandomCrop', 'RandomCropV2', 'RandomHFlip',
|
||||
'GaussianBlur', 'ColorJitter', 'RandomGray', 'ToTensor', 'Normalize',
|
||||
'ResizeRandomCrop', 'ExtractResizeRandomCrop', 'ExtractResizeAssignCrop'
|
||||
]
|
||||
|
||||
|
||||
class Compose(object):
|
||||
|
||||
def __init__(self, transforms):
|
||||
self.transforms = transforms
|
||||
|
||||
def __getitem__(self, index):
|
||||
if isinstance(index, slice):
|
||||
return Compose(self.transforms[index])
|
||||
else:
|
||||
return self.transforms[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.transforms)
|
||||
|
||||
def __call__(self, rgb):
|
||||
for t in self.transforms:
|
||||
rgb = t(rgb)
|
||||
return rgb
|
||||
|
||||
|
||||
class Resize(object):
|
||||
|
||||
def __init__(self, size=256):
|
||||
if isinstance(size, int):
|
||||
size = (size, size)
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
if isinstance(rgb, list):
|
||||
rgb = [u.resize(self.size, Image.BILINEAR) for u in rgb]
|
||||
else:
|
||||
rgb = rgb.resize(self.size, Image.BILINEAR)
|
||||
return rgb
|
||||
|
||||
|
||||
class Rescale(object):
|
||||
|
||||
def __init__(self, size=256, interpolation=Image.BILINEAR):
|
||||
self.size = size
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
scale = self.size / min(w, h)
|
||||
out_w, out_h = int(round(w * scale)), int(round(h * scale))
|
||||
rgb = [u.resize((out_w, out_h), self.interpolation) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCrop(object):
|
||||
|
||||
def __init__(self, size=224):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
assert min(w, h) >= self.size
|
||||
x1 = (w - self.size) // 2
|
||||
y1 = (h - self.size) // 2
|
||||
rgb = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ExtractResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
wh = [x1, y1, x1 + out_w, y1 + out_h]
|
||||
return rgb, wh
|
||||
|
||||
|
||||
class ExtractResizeAssignCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb, wh):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
|
||||
rgb = [u.crop(wh) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCropV2(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
# fast resize
|
||||
while min(img[0].size) >= 2 * self.size:
|
||||
img = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in img
|
||||
]
|
||||
scale = self.size / min(img[0].size)
|
||||
img = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in img
|
||||
]
|
||||
|
||||
# center crop
|
||||
x1 = (img[0].width - self.size) // 2
|
||||
y1 = (img[0].height - self.size) // 2
|
||||
img = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in img]
|
||||
return img
|
||||
|
||||
|
||||
class CenterCropWide(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
if isinstance(img, list):
|
||||
scale = min(img[0].size[0] / self.size[0],
|
||||
img[0].size[1] / self.size[1])
|
||||
img = [
|
||||
u.resize((round(u.width // scale), round(u.height // scale)),
|
||||
resample=Image.BOX) for u in img
|
||||
]
|
||||
|
||||
# center crop
|
||||
x1 = (img[0].width - self.size[0]) // 2
|
||||
y1 = (img[0].height - self.size[1]) // 2
|
||||
img = [
|
||||
u.crop((x1, y1, x1 + self.size[0], y1 + self.size[1]))
|
||||
for u in img
|
||||
]
|
||||
return img
|
||||
else:
|
||||
scale = min(img.size[0] / self.size[0], img.size[1] / self.size[1])
|
||||
img = img.resize(
|
||||
(round(img.width // scale), round(img.height // scale)),
|
||||
resample=Image.BOX)
|
||||
x1 = (img.width - self.size[0]) // 2
|
||||
y1 = (img.height - self.size[1]) // 2
|
||||
img = img.crop((x1, y1, x1 + self.size[0], y1 + self.size[1]))
|
||||
return img
|
||||
|
||||
|
||||
class RandomCrop(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4):
|
||||
self.size = size
|
||||
self.min_area = min_area
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
w, h = rgb[0].size
|
||||
area = w * h
|
||||
out_w, out_h = float('inf'), float('inf')
|
||||
while out_w > w or out_h > h:
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
aspect_ratio = random.uniform(3. / 4., 4. / 3.)
|
||||
out_w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
out_h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomCropV2(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4, ratio=(3. / 4., 4. / 3.)):
|
||||
if isinstance(size, (tuple, list)):
|
||||
self.size = size
|
||||
else:
|
||||
self.size = (size, size)
|
||||
self.min_area = min_area
|
||||
self.ratio = ratio
|
||||
|
||||
def _get_params(self, img):
|
||||
width, height = img.size
|
||||
area = height * width
|
||||
|
||||
for _ in range(10):
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
log_ratio = (math.log(self.ratio[0]), math.log(self.ratio[1]))
|
||||
aspect_ratio = math.exp(random.uniform(*log_ratio))
|
||||
|
||||
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
|
||||
if 0 < w <= width and 0 < h <= height:
|
||||
i = random.randint(0, height - h)
|
||||
j = random.randint(0, width - w)
|
||||
return i, j, h, w
|
||||
|
||||
# Fallback to central crop
|
||||
in_ratio = float(width) / float(height)
|
||||
if (in_ratio < min(self.ratio)):
|
||||
w = width
|
||||
h = int(round(w / min(self.ratio)))
|
||||
elif (in_ratio > max(self.ratio)):
|
||||
h = height
|
||||
w = int(round(h * max(self.ratio)))
|
||||
else:
|
||||
w = width
|
||||
h = height
|
||||
i = (height - h) // 2
|
||||
j = (width - w) // 2
|
||||
return i, j, h, w
|
||||
|
||||
def __call__(self, rgb):
|
||||
i, j, h, w = self._get_params(rgb[0])
|
||||
rgb = [F.resized_crop(u, i, j, h, w, self.size) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomHFlip(object):
|
||||
|
||||
def __init__(self, p=0.5):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.transpose(Image.FLIP_LEFT_RIGHT) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class GaussianBlur(object):
|
||||
|
||||
def __init__(self, sigmas=[0.1, 2.0], p=0.5):
|
||||
self.sigmas = sigmas
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
sigma = random.uniform(*self.sigmas)
|
||||
rgb = [
|
||||
u.filter(ImageFilter.GaussianBlur(radius=sigma)) for u in rgb
|
||||
]
|
||||
return rgb
|
||||
|
||||
|
||||
class ColorJitter(object):
|
||||
|
||||
def __init__(self,
|
||||
brightness=0.4,
|
||||
contrast=0.4,
|
||||
saturation=0.4,
|
||||
hue=0.1,
|
||||
p=0.5):
|
||||
self.brightness = brightness
|
||||
self.contrast = contrast
|
||||
self.saturation = saturation
|
||||
self.hue = hue
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
brightness, contrast, saturation, hue = self._random_params()
|
||||
transforms = [
|
||||
lambda f: F.adjust_brightness(f, brightness),
|
||||
lambda f: F.adjust_contrast(f, contrast),
|
||||
lambda f: F.adjust_saturation(f, saturation),
|
||||
lambda f: F.adjust_hue(f, hue)
|
||||
]
|
||||
random.shuffle(transforms)
|
||||
for t in transforms:
|
||||
rgb = [t(u) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
def _random_params(self):
|
||||
brightness = random.uniform(
|
||||
max(0, 1 - self.brightness), 1 + self.brightness)
|
||||
contrast = random.uniform(max(0, 1 - self.contrast), 1 + self.contrast)
|
||||
saturation = random.uniform(
|
||||
max(0, 1 - self.saturation), 1 + self.saturation)
|
||||
hue = random.uniform(-self.hue, self.hue)
|
||||
return brightness, contrast, saturation, hue
|
||||
|
||||
|
||||
class RandomGray(object):
|
||||
|
||||
def __init__(self, p=0.2):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.convert('L').convert('RGB') for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
|
||||
def __call__(self, rgb):
|
||||
if isinstance(rgb, list):
|
||||
rgb = torch.stack([F.to_tensor(u) for u in rgb], dim=0)
|
||||
else:
|
||||
rgb = F.to_tensor(rgb)
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
|
||||
def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
def __call__(self, rgb):
|
||||
rgb = rgb.clone()
|
||||
rgb.clamp_(0, 1)
|
||||
if not isinstance(self.mean, torch.Tensor):
|
||||
self.mean = rgb.new_tensor(self.mean).view(-1)
|
||||
if not isinstance(self.std, torch.Tensor):
|
||||
self.std = rgb.new_tensor(self.std).view(-1)
|
||||
if rgb.dim() == 4:
|
||||
rgb.sub_(self.mean.view(1, -1, 1,
|
||||
1)).div_(self.std.view(1, -1, 1, 1))
|
||||
elif rgb.dim() == 3:
|
||||
rgb.sub_(self.mean.view(-1, 1, 1)).div_(self.std.view(-1, 1, 1))
|
||||
return rgb
|
||||
228
modelscope/models/multi_modal/video_to_video/video_to_video_model.py
Executable file
228
modelscope/models/multi_modal/video_to_video/video_to_video_model.py
Executable file
@@ -0,0 +1,228 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os
|
||||
import os.path as osp
|
||||
import random
|
||||
from copy import copy
|
||||
from typing import Any, Dict
|
||||
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn.functional as F
|
||||
|
||||
import modelscope.models.multi_modal.video_to_video.utils.transforms as data
|
||||
from modelscope.metainfo import Models
|
||||
from modelscope.models.base import TorchModel
|
||||
from modelscope.models.builder import MODELS
|
||||
from modelscope.models.multi_modal.video_to_video.modules import *
|
||||
from modelscope.models.multi_modal.video_to_video.modules import (
|
||||
AutoencoderKL, FrozenOpenCLIPEmbedder, Vid2VidSDUNet,
|
||||
get_first_stage_encoding)
|
||||
from modelscope.models.multi_modal.video_to_video.utils.config import cfg
|
||||
from modelscope.models.multi_modal.video_to_video.utils.diffusion_sdedit import \
|
||||
GaussianDiffusion_SDEdit
|
||||
from modelscope.models.multi_modal.video_to_video.utils.schedules_sdedit import \
|
||||
noise_schedule
|
||||
from modelscope.models.multi_modal.video_to_video.utils.seed import setup_seed
|
||||
from modelscope.utils.config import Config
|
||||
from modelscope.utils.constant import ModelFile, Tasks
|
||||
from modelscope.utils.device import create_device
|
||||
from modelscope.utils.logger import get_logger
|
||||
|
||||
__all__ = ['VideoToVideo']
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@MODELS.register_module(
|
||||
Tasks.video_to_video, module_name=Models.video_to_video_model)
|
||||
class VideoToVideo(TorchModel):
|
||||
r"""
|
||||
Video2Video aims to solve the task of generating super-resolution videos based on input
|
||||
video and text, which is a video generation basic model developed by Alibaba Cloud.
|
||||
|
||||
Paper link: https://arxiv.org/abs/2306.02018
|
||||
|
||||
Attributes:
|
||||
diffusion: diffusion model for DDIM.
|
||||
autoencoder: decode the latent representation of input video into visual space.
|
||||
clip_encoder: encode the text into text embedding.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir, *args, **kwargs):
|
||||
r"""
|
||||
Args:
|
||||
model_dir (`str` or `os.PathLike`)
|
||||
Can be either:
|
||||
- A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co
|
||||
or modelscope.cn. Valid model ids can be located at the root-level, like `bert-base-uncased`,
|
||||
or namespaced under a user or organization name, like `dbmdz/bert-base-german-cased`.
|
||||
- A path to a *directory* containing model weights saved using
|
||||
[`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
|
||||
- A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
|
||||
this case, `from_tf` should be set to `True` and a configuration object should be provided as
|
||||
`config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
|
||||
PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
|
||||
- A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g,
|
||||
`./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to
|
||||
`True`.
|
||||
"""
|
||||
super().__init__(model_dir=model_dir, *args, **kwargs)
|
||||
|
||||
self.config = Config.from_file(
|
||||
osp.join(model_dir, ModelFile.CONFIGURATION))
|
||||
|
||||
cfg.solver_mode = self.config.model.model_args.solver_mode
|
||||
|
||||
# assign default value
|
||||
cfg.batch_size = self.config.model.model_cfg.batch_size
|
||||
cfg.target_fps = self.config.model.model_cfg.target_fps
|
||||
cfg.max_frames = self.config.model.model_cfg.max_frames
|
||||
cfg.latent_hei = self.config.model.model_cfg.latent_hei
|
||||
cfg.latent_wid = self.config.model.model_cfg.latent_wid
|
||||
cfg.model_path = osp.join(model_dir,
|
||||
self.config.model.model_args.ckpt_unet)
|
||||
|
||||
required_device = kwargs.pop('device', 'gpu')
|
||||
self.device = create_device(required_device)
|
||||
|
||||
if 'seed' in self.config.model.model_args.keys():
|
||||
cfg.seed = self.config.model.model_args.seed
|
||||
else:
|
||||
cfg.seed = random.randint(0, 99999)
|
||||
setup_seed(cfg.seed)
|
||||
|
||||
# transform
|
||||
vid_trans = data.Compose(
|
||||
[data.ToTensor(),
|
||||
data.Normalize(mean=cfg.mean, std=cfg.std)])
|
||||
self.vid_trans = vid_trans
|
||||
|
||||
cfg.embedder.pretrained = osp.join(
|
||||
model_dir, self.config.model.model_args.ckpt_clip)
|
||||
clip_encoder = FrozenOpenCLIPEmbedder(
|
||||
pretrained=cfg.embedder.pretrained, device=self.device)
|
||||
clip_encoder.model.to(self.device)
|
||||
self.clip_encoder = clip_encoder
|
||||
logger.info(f'Build encoder with {cfg.embedder.type}')
|
||||
|
||||
# [unet]
|
||||
generator = Vid2VidSDUNet()
|
||||
generator = generator.to(self.device)
|
||||
generator.eval()
|
||||
load_dict = torch.load(cfg.model_path, map_location='cpu')
|
||||
ret = generator.load_state_dict(load_dict['state_dict'], strict=True)
|
||||
self.generator = generator
|
||||
logger.info('Load model {} path {}, with local status {}'.format(
|
||||
cfg.UNet.type, cfg.model_path, ret))
|
||||
|
||||
# [diffusion]
|
||||
sigmas = noise_schedule(
|
||||
schedule='logsnr_cosine_interp',
|
||||
n=1000,
|
||||
zero_terminal_snr=True,
|
||||
scale_min=2.0,
|
||||
scale_max=4.0)
|
||||
diffusion = GaussianDiffusion_SDEdit(
|
||||
sigmas=sigmas, prediction_type='v')
|
||||
self.diffusion = diffusion
|
||||
logger.info('Build diffusion with type of GaussianDiffusion_SDEdit')
|
||||
|
||||
# [auotoencoder]
|
||||
cfg.auto_encoder.pretrained = osp.join(
|
||||
model_dir, self.config.model.model_args.ckpt_autoencoder)
|
||||
autoencoder = AutoencoderKL(**cfg.auto_encoder)
|
||||
autoencoder.eval()
|
||||
for param in autoencoder.parameters():
|
||||
param.requires_grad = False
|
||||
autoencoder.to(self.device)
|
||||
self.autoencoder = autoencoder
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
negative_prompt = cfg.negative_prompt
|
||||
negative_y = clip_encoder(negative_prompt).detach()
|
||||
self.negative_y = negative_y
|
||||
|
||||
positive_prompt = cfg.positive_prompt
|
||||
self.positive_prompt = positive_prompt
|
||||
|
||||
self.cfg = cfg
|
||||
|
||||
def forward(self, input: Dict[str, Any]):
|
||||
r"""
|
||||
The entry function of video to video task.
|
||||
1. Using CLIP to encode text into embeddings.
|
||||
2. Using diffusion model to generate the video's latent representation.
|
||||
3. Using autoencoder to decode the video's latent representation to visual space.
|
||||
|
||||
Args:
|
||||
input (`Dict[Str, Any]`):
|
||||
The input of the task
|
||||
Returns:
|
||||
A generated video (as pytorch tensor).
|
||||
"""
|
||||
|
||||
video_data = input['video_data']
|
||||
y = input['y']
|
||||
cfg = self.cfg
|
||||
|
||||
video_data = F.interpolate(
|
||||
video_data, size=(720, 1280), mode='bilinear')
|
||||
video_data = video_data.unsqueeze(0)
|
||||
video_data = video_data.to(self.device)
|
||||
|
||||
batch_size, frames_num, _, _, _ = video_data.shape
|
||||
video_data = rearrange(video_data, 'b f c h w -> (b f) c h w')
|
||||
|
||||
video_data_list = torch.chunk(
|
||||
video_data, video_data.shape[0] // 2, dim=0)
|
||||
with torch.no_grad():
|
||||
decode_data = []
|
||||
for vd_data in video_data_list:
|
||||
encoder_posterior = self.autoencoder.encode(vd_data)
|
||||
tmp = get_first_stage_encoding(encoder_posterior).detach()
|
||||
decode_data.append(tmp)
|
||||
video_data_feature = torch.cat(decode_data, dim=0)
|
||||
video_data_feature = rearrange(
|
||||
video_data_feature, '(b f) c h w -> b c f h w', b=batch_size)
|
||||
|
||||
with amp.autocast(enabled=True):
|
||||
total_noise_levels = 600
|
||||
t = torch.randint(
|
||||
total_noise_levels - 1,
|
||||
total_noise_levels, (1, ),
|
||||
dtype=torch.long).to(self.device)
|
||||
|
||||
noise = torch.randn_like(video_data_feature)
|
||||
noised_lr = self.diffusion.diffuse(video_data_feature, t, noise)
|
||||
model_kwargs = [{'y': y}, {'y': self.negative_y}]
|
||||
|
||||
gen_vid = self.diffusion.sample(
|
||||
noise=noised_lr,
|
||||
model=self.generator,
|
||||
model_kwargs=model_kwargs,
|
||||
guide_scale=7.5,
|
||||
guide_rescale=0.2,
|
||||
solver='dpmpp_2m_sde' if cfg.solver_mode == 'fast' else 'heun',
|
||||
steps=30 if cfg.solver_mode == 'fast' else 50,
|
||||
t_max=total_noise_levels - 1,
|
||||
t_min=0,
|
||||
discretization='trailing')
|
||||
|
||||
scale_factor = 0.18215
|
||||
vid_tensor_feature = 1. / scale_factor * gen_vid
|
||||
|
||||
vid_tensor_feature = rearrange(vid_tensor_feature,
|
||||
'b c f h w -> (b f) c h w')
|
||||
vid_tensor_feature_list = torch.chunk(
|
||||
vid_tensor_feature, vid_tensor_feature.shape[0] // 2, dim=0)
|
||||
decode_data = []
|
||||
for vd_data in vid_tensor_feature_list:
|
||||
tmp = self.autoencoder.decode(vd_data)
|
||||
decode_data.append(tmp)
|
||||
vid_tensor_gen = torch.cat(decode_data, dim=0)
|
||||
|
||||
gen_video = rearrange(
|
||||
vid_tensor_gen, '(b f) c h w -> b c f h w', b=cfg.batch_size)
|
||||
|
||||
return gen_video.type(torch.float32).cpu()
|
||||
23
modelscope/models/multi_modal/videocomposer/__init__.py
Normal file
23
modelscope/models/multi_modal/videocomposer/__init__.py
Normal file
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from modelscope.utils.import_utils import LazyImportModule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from .videocomposer_model import VideoComposer
|
||||
|
||||
else:
|
||||
_import_structure = {
|
||||
'videocomposer_model': ['VideoComposer'],
|
||||
}
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = LazyImportModule(
|
||||
__name__,
|
||||
globals()['__file__'],
|
||||
_import_structure,
|
||||
module_spec=__spec__,
|
||||
extra_objects={},
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from tools.annotator.util import HWC3
|
||||
|
||||
|
||||
class CannyDetector:
|
||||
|
||||
def __call__(self,
|
||||
img,
|
||||
low_threshold=None,
|
||||
high_threshold=None,
|
||||
random_threshold=True):
|
||||
|
||||
# Convert to numpy
|
||||
if isinstance(img, torch.Tensor): # (h, w, c)
|
||||
img = img.cpu().numpy()
|
||||
img_np = cv2.convertScaleAbs((img * 255.))
|
||||
elif isinstance(img, np.ndarray): # (h, w, c)
|
||||
img_np = img # we assume values are in the range from 0 to 255.
|
||||
else:
|
||||
assert False
|
||||
|
||||
# Select the threshold
|
||||
if (low_threshold is None) and (high_threshold is None):
|
||||
median_intensity = np.median(img_np)
|
||||
if random_threshold is False:
|
||||
low_threshold = int(max(0, (1 - 0.33) * median_intensity))
|
||||
high_threshold = int(min(255, (1 + 0.33) * median_intensity))
|
||||
else:
|
||||
random_canny = np.random.uniform(0.1, 0.4)
|
||||
# Might try other values
|
||||
low_threshold = int(
|
||||
max(0, (1 - random_canny) * median_intensity))
|
||||
high_threshold = 2 * low_threshold
|
||||
|
||||
# Detect canny edge
|
||||
canny_edge = cv2.Canny(img_np, low_threshold, high_threshold)
|
||||
|
||||
canny_condition = torch.from_numpy(
|
||||
canny_edge.copy()).unsqueeze(dim=-1).float().cuda() / 255.0
|
||||
return canny_condition
|
||||
@@ -0,0 +1,3 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .palette import *
|
||||
@@ -0,0 +1,155 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
r"""Modified from ``https://github.com/sergeyk/rayleigh''.
|
||||
"""
|
||||
import os
|
||||
import os.path as osp
|
||||
|
||||
import numpy as np
|
||||
from skimage.color import hsv2rgb, lab2rgb, rgb2lab
|
||||
from skimage.io import imsave
|
||||
from sklearn.metrics import euclidean_distances
|
||||
|
||||
__all__ = ['Palette']
|
||||
|
||||
|
||||
def rgb2hex(rgb):
|
||||
return '#%02x%02x%02x' % tuple([int(round(255.0 * u)) for u in rgb])
|
||||
|
||||
|
||||
def hex2rgb(hex):
|
||||
rgb = hex.strip('#')
|
||||
fn = lambda u: round(int(u, 16) / 255.0, 5) # noqa
|
||||
return fn(rgb[:2]), fn(rgb[2:4]), fn(rgb[4:6])
|
||||
|
||||
|
||||
class Palette(object):
|
||||
r"""Create a color palette (codebook) in the form of a 2D grid of colors.
|
||||
Further, the rightmost column has num_hues gradations from black to white.
|
||||
|
||||
Parameters:
|
||||
num_hues: number of colors with full lightness and saturation, in the middle.
|
||||
num_sat: number of rows above middle row that show the same hues with decreasing saturation.
|
||||
"""
|
||||
|
||||
def __init__(self, num_hues=11, num_sat=5, num_light=4):
|
||||
n = num_sat + 2 * num_light
|
||||
|
||||
# hues
|
||||
if num_hues == 8:
|
||||
hues = np.tile(
|
||||
np.array([0., 0.10, 0.15, 0.28, 0.51, 0.58, 0.77, 0.85]),
|
||||
(n, 1))
|
||||
elif num_hues == 9:
|
||||
hues = np.tile(
|
||||
np.array([0., 0.10, 0.15, 0.28, 0.49, 0.54, 0.60, 0.7, 0.87]),
|
||||
(n, 1))
|
||||
elif num_hues == 10:
|
||||
hues = np.tile(
|
||||
np.array(
|
||||
[0., 0.10, 0.15, 0.28, 0.49, 0.54, 0.60, 0.66, 0.76,
|
||||
0.87]), (n, 1))
|
||||
elif num_hues == 11:
|
||||
hues = np.tile(
|
||||
np.array([
|
||||
0.0, 0.0833, 0.166, 0.25, 0.333, 0.5, 0.56333, 0.666, 0.73,
|
||||
0.803, 0.916
|
||||
]), (n, 1))
|
||||
else:
|
||||
hues = np.tile(np.linspace(0, 1, num_hues + 1)[:-1], (n, 1))
|
||||
|
||||
# saturations
|
||||
sats = np.hstack((
|
||||
np.linspace(0, 1, num_sat + 2)[1:-1],
|
||||
1,
|
||||
[1] * num_light,
|
||||
[0.4] * # noqa
|
||||
(num_light - 1)))
|
||||
sats = np.tile(np.atleast_2d(sats).T, (1, num_hues))
|
||||
|
||||
# lights
|
||||
lights = np.hstack(
|
||||
([1] * num_sat, 1, np.linspace(1, 0.2, num_light + 2)[1:-1],
|
||||
np.linspace(1, 0.2, num_light + 2)[1:-2]))
|
||||
lights = np.tile(np.atleast_2d(lights).T, (1, num_hues))
|
||||
|
||||
# colors
|
||||
rgb = hsv2rgb(np.dstack([hues, sats, lights]))
|
||||
gray = np.tile(
|
||||
np.linspace(1, 0, n)[:, np.newaxis, np.newaxis], (1, 1, 3))
|
||||
self.thumbnail = np.hstack([rgb, gray])
|
||||
|
||||
# flatten
|
||||
rgb = rgb.T.reshape(3, -1).T
|
||||
gray = gray.T.reshape(3, -1).T
|
||||
self.rgb = np.vstack((rgb, gray))
|
||||
self.lab = rgb2lab(self.rgb[np.newaxis, :, :]).squeeze()
|
||||
self.hex = [rgb2hex(u) for u in self.rgb]
|
||||
self.lab_dists = euclidean_distances(self.lab, squared=True)
|
||||
|
||||
def histogram(self, rgb_img, sigma=20):
|
||||
# compute histogram
|
||||
lab = rgb2lab(rgb_img).reshape((-1, 3))
|
||||
min_ind = np.argmin(
|
||||
euclidean_distances(lab, self.lab, squared=True), axis=1)
|
||||
hist = 1.0 * np.bincount(
|
||||
min_ind, minlength=self.lab.shape[0]) / lab.shape[0]
|
||||
|
||||
# smooth histogram
|
||||
if sigma > 0:
|
||||
weight = np.exp(-self.lab_dists / (2.0 * sigma**2))
|
||||
weight = weight / weight.sum(1)[:, np.newaxis]
|
||||
hist = (weight * hist).sum(1)
|
||||
hist[hist < 1e-5] = 0
|
||||
return hist
|
||||
|
||||
def get_palette_image(self, hist, percentile=90, width=200, height=50):
|
||||
# curate histogram
|
||||
ind = np.argsort(-hist)
|
||||
ind = ind[hist[ind] > np.percentile(hist, percentile)]
|
||||
hist = hist[ind] / hist[ind].sum()
|
||||
|
||||
# draw palette
|
||||
nums = np.array(hist * width, dtype=int)
|
||||
array = np.vstack([
|
||||
np.tile(np.array(u), (v, 1)) for u, v in zip(self.rgb[ind], nums)
|
||||
])
|
||||
array = np.tile(array[np.newaxis, :, :], (height, 1, 1))
|
||||
if array.shape[1] < width:
|
||||
array = np.concatenate(
|
||||
[array, np.zeros((height, width - array.shape[1], 3))], axis=1)
|
||||
return array
|
||||
|
||||
def quantize_image(self, rgb_img):
|
||||
lab = rgb2lab(rgb_img).reshape((-1, 3))
|
||||
min_ind = np.argmin(
|
||||
euclidean_distances(lab, self.lab, squared=True), axis=1)
|
||||
quantized_lab = self.lab[min_ind]
|
||||
img = lab2rgb(quantized_lab.reshape(rgb_img.shape))
|
||||
return img
|
||||
|
||||
def export(self, dirname):
|
||||
if not osp.exists(dirname):
|
||||
os.makedirs(dirname)
|
||||
|
||||
# save thumbnail
|
||||
imsave(osp.join(dirname, 'palette.png'), self.thumbnail)
|
||||
|
||||
# save html
|
||||
with open(osp.join(dirname, 'palette.html'), 'w') as f:
|
||||
html = '''
|
||||
<style>
|
||||
span {
|
||||
width: 20px;
|
||||
height: 20px;
|
||||
margin: 2px;
|
||||
padding: 0px;
|
||||
display: inline-block;
|
||||
}
|
||||
</style>
|
||||
'''
|
||||
for row in self.thumbnail:
|
||||
for col in row:
|
||||
html += '<a id="{0}"><span style="background-color: {0}" /></a>\n'.format(
|
||||
rgb2hex(col))
|
||||
html += '<br />\n'
|
||||
f.write(html)
|
||||
@@ -0,0 +1,4 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .pidinet import *
|
||||
from .sketch_simplification import *
|
||||
@@ -0,0 +1,940 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
r"""Modified from ``https://github.com/zhuoinoulu/pidinet''.
|
||||
Image augmentation: T.Compose([
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]).
|
||||
"""
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modelscope.models.multi_modal.videocomposer.utils.utils import \
|
||||
DOWNLOAD_TO_CACHE
|
||||
|
||||
__all__ = [
|
||||
'PiDiNet', 'pidinet_bsd_tiny', 'pidinet_bsd_small', 'pidinet_bsd',
|
||||
'pidinet_nyud', 'pidinet_multicue'
|
||||
]
|
||||
|
||||
CONFIGS = {
|
||||
'baseline': {
|
||||
'layer0': 'cv',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cv',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cv',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cv',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'c-v15': {
|
||||
'layer0': 'cd',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cv',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cv',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cv',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'a-v15': {
|
||||
'layer0': 'ad',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cv',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cv',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cv',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'r-v15': {
|
||||
'layer0': 'rd',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cv',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cv',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cv',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'cvvv4': {
|
||||
'layer0': 'cd',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cd',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cd',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cd',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'avvv4': {
|
||||
'layer0': 'ad',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'ad',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'ad',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'ad',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'rvvv4': {
|
||||
'layer0': 'rd',
|
||||
'layer1': 'cv',
|
||||
'layer2': 'cv',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'rd',
|
||||
'layer5': 'cv',
|
||||
'layer6': 'cv',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'rd',
|
||||
'layer9': 'cv',
|
||||
'layer10': 'cv',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'rd',
|
||||
'layer13': 'cv',
|
||||
'layer14': 'cv',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'cccv4': {
|
||||
'layer0': 'cd',
|
||||
'layer1': 'cd',
|
||||
'layer2': 'cd',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cd',
|
||||
'layer5': 'cd',
|
||||
'layer6': 'cd',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cd',
|
||||
'layer9': 'cd',
|
||||
'layer10': 'cd',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cd',
|
||||
'layer13': 'cd',
|
||||
'layer14': 'cd',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'aaav4': {
|
||||
'layer0': 'ad',
|
||||
'layer1': 'ad',
|
||||
'layer2': 'ad',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'ad',
|
||||
'layer5': 'ad',
|
||||
'layer6': 'ad',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'ad',
|
||||
'layer9': 'ad',
|
||||
'layer10': 'ad',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'ad',
|
||||
'layer13': 'ad',
|
||||
'layer14': 'ad',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'rrrv4': {
|
||||
'layer0': 'rd',
|
||||
'layer1': 'rd',
|
||||
'layer2': 'rd',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'rd',
|
||||
'layer5': 'rd',
|
||||
'layer6': 'rd',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'rd',
|
||||
'layer9': 'rd',
|
||||
'layer10': 'rd',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'rd',
|
||||
'layer13': 'rd',
|
||||
'layer14': 'rd',
|
||||
'layer15': 'cv',
|
||||
},
|
||||
'c16': {
|
||||
'layer0': 'cd',
|
||||
'layer1': 'cd',
|
||||
'layer2': 'cd',
|
||||
'layer3': 'cd',
|
||||
'layer4': 'cd',
|
||||
'layer5': 'cd',
|
||||
'layer6': 'cd',
|
||||
'layer7': 'cd',
|
||||
'layer8': 'cd',
|
||||
'layer9': 'cd',
|
||||
'layer10': 'cd',
|
||||
'layer11': 'cd',
|
||||
'layer12': 'cd',
|
||||
'layer13': 'cd',
|
||||
'layer14': 'cd',
|
||||
'layer15': 'cd',
|
||||
},
|
||||
'a16': {
|
||||
'layer0': 'ad',
|
||||
'layer1': 'ad',
|
||||
'layer2': 'ad',
|
||||
'layer3': 'ad',
|
||||
'layer4': 'ad',
|
||||
'layer5': 'ad',
|
||||
'layer6': 'ad',
|
||||
'layer7': 'ad',
|
||||
'layer8': 'ad',
|
||||
'layer9': 'ad',
|
||||
'layer10': 'ad',
|
||||
'layer11': 'ad',
|
||||
'layer12': 'ad',
|
||||
'layer13': 'ad',
|
||||
'layer14': 'ad',
|
||||
'layer15': 'ad',
|
||||
},
|
||||
'r16': {
|
||||
'layer0': 'rd',
|
||||
'layer1': 'rd',
|
||||
'layer2': 'rd',
|
||||
'layer3': 'rd',
|
||||
'layer4': 'rd',
|
||||
'layer5': 'rd',
|
||||
'layer6': 'rd',
|
||||
'layer7': 'rd',
|
||||
'layer8': 'rd',
|
||||
'layer9': 'rd',
|
||||
'layer10': 'rd',
|
||||
'layer11': 'rd',
|
||||
'layer12': 'rd',
|
||||
'layer13': 'rd',
|
||||
'layer14': 'rd',
|
||||
'layer15': 'rd',
|
||||
},
|
||||
'carv4': {
|
||||
'layer0': 'cd',
|
||||
'layer1': 'ad',
|
||||
'layer2': 'rd',
|
||||
'layer3': 'cv',
|
||||
'layer4': 'cd',
|
||||
'layer5': 'ad',
|
||||
'layer6': 'rd',
|
||||
'layer7': 'cv',
|
||||
'layer8': 'cd',
|
||||
'layer9': 'ad',
|
||||
'layer10': 'rd',
|
||||
'layer11': 'cv',
|
||||
'layer12': 'cd',
|
||||
'layer13': 'ad',
|
||||
'layer14': 'rd',
|
||||
'layer15': 'cv'
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def create_conv_func(op_type):
|
||||
assert op_type in ['cv', 'cd', 'ad',
|
||||
'rd'], 'unknown op type: %s' % str(op_type)
|
||||
if op_type == 'cv':
|
||||
return F.conv2d
|
||||
if op_type == 'cd':
|
||||
|
||||
def func(x,
|
||||
weights,
|
||||
bias=None,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1):
|
||||
assert dilation in [1,
|
||||
2], 'dilation for cd_conv should be in 1 or 2'
|
||||
assert weights.size(2) == 3 and weights.size(
|
||||
3) == 3, 'kernel size for cd_conv should be 3x3'
|
||||
assert padding == dilation, 'padding for cd_conv set wrong'
|
||||
|
||||
weights_c = weights.sum(dim=[2, 3], keepdim=True)
|
||||
yc = F.conv2d(
|
||||
x, weights_c, stride=stride, padding=0, groups=groups)
|
||||
y = F.conv2d(
|
||||
x,
|
||||
weights,
|
||||
bias,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups)
|
||||
return y - yc
|
||||
|
||||
return func
|
||||
elif op_type == 'ad':
|
||||
|
||||
def func(x,
|
||||
weights,
|
||||
bias=None,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1):
|
||||
assert dilation in [1,
|
||||
2], 'dilation for ad_conv should be in 1 or 2'
|
||||
assert weights.size(2) == 3 and weights.size(
|
||||
3) == 3, 'kernel size for ad_conv should be 3x3'
|
||||
assert padding == dilation, 'padding for ad_conv set wrong'
|
||||
|
||||
shape = weights.shape
|
||||
weights = weights.view(shape[0], shape[1], -1)
|
||||
weights_conv = (weights
|
||||
- weights[:, :, [3, 0, 1, 6, 4, 2, 7, 8, 5]]).view(
|
||||
shape) # clock-wise
|
||||
y = F.conv2d(
|
||||
x,
|
||||
weights_conv,
|
||||
bias,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups)
|
||||
return y
|
||||
|
||||
return func
|
||||
elif op_type == 'rd':
|
||||
|
||||
def func(x,
|
||||
weights,
|
||||
bias=None,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1):
|
||||
assert dilation in [1,
|
||||
2], 'dilation for rd_conv should be in 1 or 2'
|
||||
assert weights.size(2) == 3 and weights.size(
|
||||
3) == 3, 'kernel size for rd_conv should be 3x3'
|
||||
padding = 2 * dilation
|
||||
|
||||
shape = weights.shape
|
||||
if weights.is_cuda:
|
||||
buffer = torch.cuda.FloatTensor(shape[0], shape[1],
|
||||
5 * 5).fill_(0)
|
||||
else:
|
||||
buffer = torch.zeros(shape[0], shape[1], 5 * 5)
|
||||
weights = weights.view(shape[0], shape[1], -1)
|
||||
buffer[:, :, [0, 2, 4, 10, 14, 20, 22, 24]] = weights[:, :, 1:]
|
||||
buffer[:, :, [6, 7, 8, 11, 13, 16, 17, 18]] = -weights[:, :, 1:]
|
||||
buffer[:, :, 12] = 0
|
||||
buffer = buffer.view(shape[0], shape[1], 5, 5)
|
||||
y = F.conv2d(
|
||||
x,
|
||||
buffer,
|
||||
bias,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups)
|
||||
return y
|
||||
|
||||
return func
|
||||
else:
|
||||
print('impossible to be here unless you force that', flush=True)
|
||||
return None
|
||||
|
||||
|
||||
def config_model(model):
|
||||
model_options = list(CONFIGS.keys())
|
||||
assert model in model_options, \
|
||||
'unrecognized model, please choose from %s' % str(model_options)
|
||||
|
||||
pdcs = []
|
||||
for i in range(16):
|
||||
layer_name = 'layer%d' % i
|
||||
op = CONFIGS[model][layer_name]
|
||||
pdcs.append(create_conv_func(op))
|
||||
return pdcs
|
||||
|
||||
|
||||
def config_model_converted(model):
|
||||
model_options = list(CONFIGS.keys())
|
||||
assert model in model_options, \
|
||||
'unrecognized model, please choose from %s' % str(model_options)
|
||||
|
||||
pdcs = []
|
||||
for i in range(16):
|
||||
layer_name = 'layer%d' % i
|
||||
op = CONFIGS[model][layer_name]
|
||||
pdcs.append(op)
|
||||
return pdcs
|
||||
|
||||
|
||||
def convert_pdc(op, weight):
|
||||
if op == 'cv':
|
||||
return weight
|
||||
elif op == 'cd':
|
||||
shape = weight.shape
|
||||
weight_c = weight.sum(dim=[2, 3])
|
||||
weight = weight.view(shape[0], shape[1], -1)
|
||||
weight[:, :, 4] = weight[:, :, 4] - weight_c
|
||||
weight = weight.view(shape)
|
||||
return weight
|
||||
elif op == 'ad':
|
||||
shape = weight.shape
|
||||
weight = weight.view(shape[0], shape[1], -1)
|
||||
weight_conv = (weight
|
||||
- weight[:, :, [3, 0, 1, 6, 4, 2, 7, 8, 5]]).view(shape)
|
||||
return weight_conv
|
||||
elif op == 'rd':
|
||||
shape = weight.shape
|
||||
buffer = torch.zeros(shape[0], shape[1], 5 * 5, device=weight.device)
|
||||
weight = weight.view(shape[0], shape[1], -1)
|
||||
buffer[:, :, [0, 2, 4, 10, 14, 20, 22, 24]] = weight[:, :, 1:]
|
||||
buffer[:, :, [6, 7, 8, 11, 13, 16, 17, 18]] = -weight[:, :, 1:]
|
||||
buffer = buffer.view(shape[0], shape[1], 5, 5)
|
||||
return buffer
|
||||
raise ValueError('wrong op {}'.format(str(op)))
|
||||
|
||||
|
||||
def convert_pidinet(state_dict, config):
|
||||
pdcs = config_model_converted(config)
|
||||
new_dict = {}
|
||||
for pname, p in state_dict.items():
|
||||
if 'init_block.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[0], p)
|
||||
elif 'block1_1.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[1], p)
|
||||
elif 'block1_2.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[2], p)
|
||||
elif 'block1_3.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[3], p)
|
||||
elif 'block2_1.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[4], p)
|
||||
elif 'block2_2.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[5], p)
|
||||
elif 'block2_3.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[6], p)
|
||||
elif 'block2_4.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[7], p)
|
||||
elif 'block3_1.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[8], p)
|
||||
elif 'block3_2.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[9], p)
|
||||
elif 'block3_3.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[10], p)
|
||||
elif 'block3_4.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[11], p)
|
||||
elif 'block4_1.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[12], p)
|
||||
elif 'block4_2.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[13], p)
|
||||
elif 'block4_3.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[14], p)
|
||||
elif 'block4_4.conv1.weight' in pname:
|
||||
new_dict[pname] = convert_pdc(pdcs[15], p)
|
||||
else:
|
||||
new_dict[pname] = p
|
||||
return new_dict
|
||||
|
||||
|
||||
class Conv2d(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
pdc,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=False):
|
||||
super(Conv2d, self).__init__()
|
||||
if in_channels % groups != 0:
|
||||
raise ValueError('in_channels must be divisible by groups')
|
||||
if out_channels % groups != 0:
|
||||
raise ValueError('out_channels must be divisible by groups')
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.dilation = dilation
|
||||
self.groups = groups
|
||||
self.weight = nn.Parameter(
|
||||
torch.Tensor(out_channels, in_channels // groups, kernel_size,
|
||||
kernel_size))
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.Tensor(out_channels))
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
self.reset_parameters()
|
||||
self.pdc = pdc
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
nn.init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, input):
|
||||
return self.pdc(input, self.weight, self.bias, self.stride,
|
||||
self.padding, self.dilation, self.groups)
|
||||
|
||||
|
||||
class CSAM(nn.Module):
|
||||
r"""
|
||||
Compact Spatial Attention Module
|
||||
"""
|
||||
|
||||
def __init__(self, channels):
|
||||
super(CSAM, self).__init__()
|
||||
|
||||
mid_channels = 4
|
||||
self.relu1 = nn.ReLU()
|
||||
self.conv1 = nn.Conv2d(
|
||||
channels, mid_channels, kernel_size=1, padding=0)
|
||||
self.conv2 = nn.Conv2d(
|
||||
mid_channels, 1, kernel_size=3, padding=1, bias=False)
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
nn.init.constant_(self.conv1.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.relu1(x)
|
||||
y = self.conv1(y)
|
||||
y = self.conv2(y)
|
||||
y = self.sigmoid(y)
|
||||
|
||||
return x * y
|
||||
|
||||
|
||||
class CDCM(nn.Module):
|
||||
r"""
|
||||
Compact Dilation Convolution based Module
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super(CDCM, self).__init__()
|
||||
|
||||
self.relu1 = nn.ReLU()
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, padding=0)
|
||||
self.conv2_1 = nn.Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
dilation=5,
|
||||
padding=5,
|
||||
bias=False)
|
||||
self.conv2_2 = nn.Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
dilation=7,
|
||||
padding=7,
|
||||
bias=False)
|
||||
self.conv2_3 = nn.Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
dilation=9,
|
||||
padding=9,
|
||||
bias=False)
|
||||
self.conv2_4 = nn.Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
dilation=11,
|
||||
padding=11,
|
||||
bias=False)
|
||||
nn.init.constant_(self.conv1.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.relu1(x)
|
||||
x = self.conv1(x)
|
||||
x1 = self.conv2_1(x)
|
||||
x2 = self.conv2_2(x)
|
||||
x3 = self.conv2_3(x)
|
||||
x4 = self.conv2_4(x)
|
||||
return x1 + x2 + x3 + x4
|
||||
|
||||
|
||||
class MapReduce(nn.Module):
|
||||
r"""
|
||||
Reduce feature maps into a single edge map
|
||||
"""
|
||||
|
||||
def __init__(self, channels):
|
||||
super(MapReduce, self).__init__()
|
||||
self.conv = nn.Conv2d(channels, 1, kernel_size=1, padding=0)
|
||||
nn.init.constant_(self.conv.bias, 0)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class PDCBlock(nn.Module):
|
||||
|
||||
def __init__(self, pdc, inplane, ouplane, stride=1):
|
||||
super(PDCBlock, self).__init__()
|
||||
self.stride = stride
|
||||
|
||||
self.stride = stride
|
||||
if self.stride > 1:
|
||||
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
self.shortcut = nn.Conv2d(
|
||||
inplane, ouplane, kernel_size=1, padding=0)
|
||||
self.conv1 = Conv2d(
|
||||
pdc,
|
||||
inplane,
|
||||
inplane,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
groups=inplane,
|
||||
bias=False)
|
||||
self.relu2 = nn.ReLU()
|
||||
self.conv2 = nn.Conv2d(
|
||||
inplane, ouplane, kernel_size=1, padding=0, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
if self.stride > 1:
|
||||
x = self.pool(x)
|
||||
y = self.conv1(x)
|
||||
y = self.relu2(y)
|
||||
y = self.conv2(y)
|
||||
if self.stride > 1:
|
||||
x = self.shortcut(x)
|
||||
y = y + x
|
||||
return y
|
||||
|
||||
|
||||
class PDCBlock_converted(nn.Module):
|
||||
r"""
|
||||
CPDC, APDC can be converted to vanilla 3x3 convolution
|
||||
RPDC can be converted to vanilla 5x5 convolution
|
||||
"""
|
||||
|
||||
def __init__(self, pdc, inplane, ouplane, stride=1):
|
||||
super(PDCBlock_converted, self).__init__()
|
||||
self.stride = stride
|
||||
|
||||
if self.stride > 1:
|
||||
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
self.shortcut = nn.Conv2d(
|
||||
inplane, ouplane, kernel_size=1, padding=0)
|
||||
if pdc == 'rd':
|
||||
self.conv1 = nn.Conv2d(
|
||||
inplane,
|
||||
inplane,
|
||||
kernel_size=5,
|
||||
padding=2,
|
||||
groups=inplane,
|
||||
bias=False)
|
||||
else:
|
||||
self.conv1 = nn.Conv2d(
|
||||
inplane,
|
||||
inplane,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
groups=inplane,
|
||||
bias=False)
|
||||
self.relu2 = nn.ReLU()
|
||||
self.conv2 = nn.Conv2d(
|
||||
inplane, ouplane, kernel_size=1, padding=0, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
if self.stride > 1:
|
||||
x = self.pool(x)
|
||||
y = self.conv1(x)
|
||||
y = self.relu2(y)
|
||||
y = self.conv2(y)
|
||||
if self.stride > 1:
|
||||
x = self.shortcut(x)
|
||||
y = y + x
|
||||
return y
|
||||
|
||||
|
||||
class PiDiNet(nn.Module):
|
||||
|
||||
def __init__(self, inplane, pdcs, dil=None, sa=False, convert=False):
|
||||
super(PiDiNet, self).__init__()
|
||||
self.sa = sa
|
||||
if dil is not None:
|
||||
assert isinstance(dil, int), 'dil should be an int'
|
||||
self.dil = dil
|
||||
|
||||
self.fuseplanes = []
|
||||
|
||||
self.inplane = inplane
|
||||
if convert:
|
||||
if pdcs[0] == 'rd':
|
||||
init_kernel_size = 5
|
||||
init_padding = 2
|
||||
else:
|
||||
init_kernel_size = 3
|
||||
init_padding = 1
|
||||
self.init_block = nn.Conv2d(
|
||||
3,
|
||||
self.inplane,
|
||||
kernel_size=init_kernel_size,
|
||||
padding=init_padding,
|
||||
bias=False)
|
||||
block_class = PDCBlock_converted
|
||||
else:
|
||||
self.init_block = Conv2d(
|
||||
pdcs[0], 3, self.inplane, kernel_size=3, padding=1)
|
||||
block_class = PDCBlock
|
||||
|
||||
self.block1_1 = block_class(pdcs[1], self.inplane, self.inplane)
|
||||
self.block1_2 = block_class(pdcs[2], self.inplane, self.inplane)
|
||||
self.block1_3 = block_class(pdcs[3], self.inplane, self.inplane)
|
||||
self.fuseplanes.append(self.inplane) # C
|
||||
|
||||
inplane = self.inplane
|
||||
self.inplane = self.inplane * 2
|
||||
self.block2_1 = block_class(pdcs[4], inplane, self.inplane, stride=2)
|
||||
self.block2_2 = block_class(pdcs[5], self.inplane, self.inplane)
|
||||
self.block2_3 = block_class(pdcs[6], self.inplane, self.inplane)
|
||||
self.block2_4 = block_class(pdcs[7], self.inplane, self.inplane)
|
||||
self.fuseplanes.append(self.inplane) # 2C
|
||||
|
||||
inplane = self.inplane
|
||||
self.inplane = self.inplane * 2
|
||||
self.block3_1 = block_class(pdcs[8], inplane, self.inplane, stride=2)
|
||||
self.block3_2 = block_class(pdcs[9], self.inplane, self.inplane)
|
||||
self.block3_3 = block_class(pdcs[10], self.inplane, self.inplane)
|
||||
self.block3_4 = block_class(pdcs[11], self.inplane, self.inplane)
|
||||
self.fuseplanes.append(self.inplane) # 4C
|
||||
|
||||
self.block4_1 = block_class(
|
||||
pdcs[12], self.inplane, self.inplane, stride=2)
|
||||
self.block4_2 = block_class(pdcs[13], self.inplane, self.inplane)
|
||||
self.block4_3 = block_class(pdcs[14], self.inplane, self.inplane)
|
||||
self.block4_4 = block_class(pdcs[15], self.inplane, self.inplane)
|
||||
self.fuseplanes.append(self.inplane) # 4C
|
||||
|
||||
self.conv_reduces = nn.ModuleList()
|
||||
if self.sa and self.dil is not None:
|
||||
self.attentions = nn.ModuleList()
|
||||
self.dilations = nn.ModuleList()
|
||||
for i in range(4):
|
||||
self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
|
||||
self.attentions.append(CSAM(self.dil))
|
||||
self.conv_reduces.append(MapReduce(self.dil))
|
||||
elif self.sa:
|
||||
self.attentions = nn.ModuleList()
|
||||
for i in range(4):
|
||||
self.attentions.append(CSAM(self.fuseplanes[i]))
|
||||
self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
|
||||
elif self.dil is not None:
|
||||
self.dilations = nn.ModuleList()
|
||||
for i in range(4):
|
||||
self.dilations.append(CDCM(self.fuseplanes[i], self.dil))
|
||||
self.conv_reduces.append(MapReduce(self.dil))
|
||||
else:
|
||||
for i in range(4):
|
||||
self.conv_reduces.append(MapReduce(self.fuseplanes[i]))
|
||||
|
||||
self.classifier = nn.Conv2d(4, 1, kernel_size=1)
|
||||
nn.init.constant_(self.classifier.weight, 0.25)
|
||||
nn.init.constant_(self.classifier.bias, 0)
|
||||
|
||||
def get_weights(self):
|
||||
conv_weights = []
|
||||
bn_weights = []
|
||||
relu_weights = []
|
||||
for pname, p in self.named_parameters():
|
||||
if 'bn' in pname:
|
||||
bn_weights.append(p)
|
||||
elif 'relu' in pname:
|
||||
relu_weights.append(p)
|
||||
else:
|
||||
conv_weights.append(p)
|
||||
|
||||
return conv_weights, bn_weights, relu_weights
|
||||
|
||||
def forward(self, x):
|
||||
H, W = x.size()[2:]
|
||||
|
||||
x = self.init_block(x)
|
||||
|
||||
x1 = self.block1_1(x)
|
||||
x1 = self.block1_2(x1)
|
||||
x1 = self.block1_3(x1)
|
||||
|
||||
x2 = self.block2_1(x1)
|
||||
x2 = self.block2_2(x2)
|
||||
x2 = self.block2_3(x2)
|
||||
x2 = self.block2_4(x2)
|
||||
|
||||
x3 = self.block3_1(x2)
|
||||
x3 = self.block3_2(x3)
|
||||
x3 = self.block3_3(x3)
|
||||
x3 = self.block3_4(x3)
|
||||
|
||||
x4 = self.block4_1(x3)
|
||||
x4 = self.block4_2(x4)
|
||||
x4 = self.block4_3(x4)
|
||||
x4 = self.block4_4(x4)
|
||||
|
||||
x_fuses = []
|
||||
if self.sa and self.dil is not None:
|
||||
for i, xi in enumerate([x1, x2, x3, x4]):
|
||||
x_fuses.append(self.attentions[i](self.dilations[i](xi)))
|
||||
elif self.sa:
|
||||
for i, xi in enumerate([x1, x2, x3, x4]):
|
||||
x_fuses.append(self.attentions[i](xi))
|
||||
elif self.dil is not None:
|
||||
for i, xi in enumerate([x1, x2, x3, x4]):
|
||||
x_fuses.append(self.dilations[i](xi))
|
||||
else:
|
||||
x_fuses = [x1, x2, x3, x4]
|
||||
|
||||
e1 = self.conv_reduces[0](x_fuses[0])
|
||||
e1 = F.interpolate(e1, (H, W), mode='bilinear', align_corners=False)
|
||||
|
||||
e2 = self.conv_reduces[1](x_fuses[1])
|
||||
e2 = F.interpolate(e2, (H, W), mode='bilinear', align_corners=False)
|
||||
|
||||
e3 = self.conv_reduces[2](x_fuses[2])
|
||||
e3 = F.interpolate(e3, (H, W), mode='bilinear', align_corners=False)
|
||||
|
||||
e4 = self.conv_reduces[3](x_fuses[3])
|
||||
e4 = F.interpolate(e4, (H, W), mode='bilinear', align_corners=False)
|
||||
|
||||
outputs = [e1, e2, e3, e4]
|
||||
output = self.classifier(torch.cat(outputs, dim=1))
|
||||
|
||||
outputs.append(output)
|
||||
outputs = [torch.sigmoid(r) for r in outputs]
|
||||
return outputs[-1]
|
||||
|
||||
|
||||
def pidinet_bsd_tiny(pretrained=False, vanilla_cnn=True):
|
||||
pdcs = config_model_converted('carv4') if vanilla_cnn else config_model(
|
||||
'carv4')
|
||||
model = PiDiNet(20, pdcs, dil=8, sa=True, convert=vanilla_cnn)
|
||||
if pretrained:
|
||||
state = torch.load(
|
||||
DOWNLOAD_TO_CACHE('models/pidinet/table5_pidinet-tiny.pth'),
|
||||
map_location='cpu')['state_dict']
|
||||
if vanilla_cnn:
|
||||
state = convert_pidinet(state, 'carv4')
|
||||
state = {
|
||||
k[len('module.'):] if k.startswith('module.') else k: v
|
||||
for k, v in state.items()
|
||||
}
|
||||
model.load_state_dict(state)
|
||||
return model
|
||||
|
||||
|
||||
def pidinet_bsd_small(pretrained=False, vanilla_cnn=True):
|
||||
pdcs = config_model_converted('carv4') if vanilla_cnn else config_model(
|
||||
'carv4')
|
||||
model = PiDiNet(30, pdcs, dil=12, sa=True, convert=vanilla_cnn)
|
||||
if pretrained:
|
||||
state = torch.load(
|
||||
DOWNLOAD_TO_CACHE('models/pidinet/table5_pidinet-small.pth'),
|
||||
map_location='cpu')['state_dict']
|
||||
if vanilla_cnn:
|
||||
state = convert_pidinet(state, 'carv4')
|
||||
state = {
|
||||
k[len('module.'):] if k.startswith('module.') else k: v
|
||||
for k, v in state.items()
|
||||
}
|
||||
model.load_state_dict(state)
|
||||
return model
|
||||
|
||||
|
||||
def pidinet_bsd(model_dir, pretrained=False, vanilla_cnn=True):
|
||||
pdcs = config_model_converted('carv4') if vanilla_cnn else config_model(
|
||||
'carv4')
|
||||
model = PiDiNet(60, pdcs, dil=24, sa=True, convert=vanilla_cnn)
|
||||
if pretrained:
|
||||
state = torch.load(
|
||||
os.path.join(model_dir, 'table5_pidinet.pth'),
|
||||
map_location='cpu')['state_dict']
|
||||
if vanilla_cnn:
|
||||
state = convert_pidinet(state, 'carv4')
|
||||
state = {
|
||||
k[len('module.'):] if k.startswith('module.') else k: v
|
||||
for k, v in state.items()
|
||||
}
|
||||
model.load_state_dict(state)
|
||||
return model
|
||||
|
||||
|
||||
def pidinet_nyud(pretrained=False, vanilla_cnn=True):
|
||||
pdcs = config_model_converted('carv4') if vanilla_cnn else config_model(
|
||||
'carv4')
|
||||
model = PiDiNet(60, pdcs, dil=24, sa=True, convert=vanilla_cnn)
|
||||
if pretrained:
|
||||
state = torch.load(
|
||||
DOWNLOAD_TO_CACHE('models/pidinet/table6_pidinet.pth'),
|
||||
map_location='cpu')['state_dict']
|
||||
if vanilla_cnn:
|
||||
state = convert_pidinet(state, 'carv4')
|
||||
state = {
|
||||
k[len('module.'):] if k.startswith('module.') else k: v
|
||||
for k, v in state.items()
|
||||
}
|
||||
model.load_state_dict(state)
|
||||
return model
|
||||
|
||||
|
||||
def pidinet_multicue(pretrained=False, vanilla_cnn=True):
|
||||
pdcs = config_model_converted('carv4') if vanilla_cnn else config_model(
|
||||
'carv4')
|
||||
model = PiDiNet(60, pdcs, dil=24, sa=True, convert=vanilla_cnn)
|
||||
if pretrained:
|
||||
state = torch.load(
|
||||
DOWNLOAD_TO_CACHE('models/pidinet/table7_pidinet.pth'),
|
||||
map_location='cpu')['state_dict']
|
||||
if vanilla_cnn:
|
||||
state = convert_pidinet(state, 'carv4')
|
||||
state = {
|
||||
k[len('module.'):] if k.startswith('module.') else k: v
|
||||
for k, v in state.items()
|
||||
}
|
||||
model.load_state_dict(state)
|
||||
return model
|
||||
@@ -0,0 +1,110 @@
|
||||
r"""PyTorch re-implementation adapted from the Lua code in ``https://github.com/bobbens/sketch_simplification''.
|
||||
"""
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modelscope.models.multi_modal.videocomposer.utils.utils import \
|
||||
DOWNLOAD_TO_CACHE
|
||||
|
||||
__all__ = [
|
||||
'SketchSimplification', 'sketch_simplification_gan',
|
||||
'sketch_simplification_mse', 'sketch_to_pencil_v1', 'sketch_to_pencil_v2'
|
||||
]
|
||||
|
||||
|
||||
class SketchSimplification(nn.Module):
|
||||
r"""NOTE:
|
||||
1. Input image should has only one gray channel.
|
||||
2. Input image size should be divisible by 8.
|
||||
3. Sketch in the input/output image is in dark color while background in light color.
|
||||
"""
|
||||
|
||||
def __init__(self, mean, std):
|
||||
assert isinstance(mean, float) and isinstance(std, float)
|
||||
super(SketchSimplification, self).__init__()
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
# layers
|
||||
self.layers = nn.Sequential(
|
||||
nn.Conv2d(1, 48, 5, 2, 2), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(48, 128, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(128, 128, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(128, 128, 3, 2, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(128, 256, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 256, 3, 2, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 512, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(512, 1024, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(1024, 1024, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(1024, 1024, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(1024, 1024, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(1024, 512, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(512, 256, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.ConvTranspose2d(256, 256, 4, 2, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 128, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.ConvTranspose2d(128, 128, 4, 2, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(128, 128, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(128, 48, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.ConvTranspose2d(48, 48, 4, 2, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(48, 24, 3, 1, 1), nn.ReLU(inplace=True),
|
||||
nn.Conv2d(24, 1, 3, 1, 1), nn.Sigmoid())
|
||||
|
||||
def forward(self, x):
|
||||
r"""x: [B, 1, H, W] within range [0, 1]. Sketch pixels in dark color.
|
||||
"""
|
||||
x = (x - self.mean) / self.std
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
def sketch_simplification_gan(model_dir, pretrained=False):
|
||||
model = SketchSimplification(
|
||||
mean=0.9664114577640158, std=0.0858381272736797)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(model_dir, 'sketch_simplification_gan.pth'),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
|
||||
|
||||
def sketch_simplification_mse(pretrained=False):
|
||||
model = SketchSimplification(
|
||||
mean=0.9664423107454593, std=0.08583666033640507)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
DOWNLOAD_TO_CACHE(
|
||||
'models/sketch_simplification/sketch_simplification_mse.pth'
|
||||
),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
|
||||
|
||||
def sketch_to_pencil_v1(pretrained=False):
|
||||
model = SketchSimplification(
|
||||
mean=0.9817833515894078, std=0.0925009022585048)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
DOWNLOAD_TO_CACHE(
|
||||
'models/sketch_simplification/sketch_to_pencil_v1.pth'),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
|
||||
|
||||
def sketch_to_pencil_v2(pretrained=False):
|
||||
model = SketchSimplification(
|
||||
mean=0.9851298627337799, std=0.07418377454883571)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
DOWNLOAD_TO_CACHE(
|
||||
'models/sketch_simplification/sketch_to_pencil_v2.pth'),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
@@ -0,0 +1,42 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')
|
||||
|
||||
|
||||
def HWC3(x):
|
||||
assert x.dtype == np.uint8
|
||||
if x.ndim == 2:
|
||||
x = x[:, :, None]
|
||||
assert x.ndim == 3
|
||||
H, W, C = x.shape
|
||||
assert C == 1 or C == 3 or C == 4
|
||||
if C == 3:
|
||||
return x
|
||||
if C == 1:
|
||||
return np.concatenate([x, x, x], axis=2)
|
||||
if C == 4:
|
||||
color = x[:, :, 0:3].astype(np.float32)
|
||||
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
||||
y = color * alpha + 255.0 * (1.0 - alpha)
|
||||
y = y.clip(0, 255).astype(np.uint8)
|
||||
return y
|
||||
|
||||
|
||||
def resize_image(input_image, resolution):
|
||||
H, W, C = input_image.shape
|
||||
H = float(H)
|
||||
W = float(W)
|
||||
k = float(resolution) / min(H, W)
|
||||
H *= k
|
||||
W *= k
|
||||
H = int(np.round(H / 64.0)) * 64
|
||||
W = int(np.round(W / 64.0)) * 64
|
||||
img = cv2.resize(
|
||||
input_image, (W, H),
|
||||
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
||||
return img
|
||||
650
modelscope/models/multi_modal/videocomposer/autoencoder.py
Normal file
650
modelscope/models/multi_modal/videocomposer/autoencoder.py
Normal file
@@ -0,0 +1,650 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
__all__ = ['AutoencoderKL']
|
||||
|
||||
|
||||
def nonlinearity(x):
|
||||
# swish
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
def Normalize(in_channels, num_groups=32):
|
||||
return torch.nn.GroupNorm(
|
||||
num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
|
||||
class DiagonalGaussianDistribution(object):
|
||||
|
||||
def __init__(self, parameters, deterministic=False):
|
||||
self.parameters = parameters
|
||||
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
|
||||
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
|
||||
self.deterministic = deterministic
|
||||
self.std = torch.exp(0.5 * self.logvar)
|
||||
self.var = torch.exp(self.logvar)
|
||||
if self.deterministic:
|
||||
self.var = self.std = torch.zeros_like(
|
||||
self.mean).to(device=self.parameters.device)
|
||||
|
||||
def sample(self):
|
||||
x = self.mean + self.std * torch.randn(
|
||||
self.mean.shape).to(device=self.parameters.device)
|
||||
return x
|
||||
|
||||
def kl(self, other=None):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
else:
|
||||
if other is None:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
|
||||
dim=[1, 2, 3])
|
||||
else:
|
||||
return 0.5 * torch.sum(
|
||||
torch.pow(self.mean - other.mean, 2) / other.var
|
||||
+ self.var / other.var - 1.0 - self.logvar + other.logvar,
|
||||
dim=[1, 2, 3])
|
||||
|
||||
def nll(self, sample, dims=[1, 2, 3]):
|
||||
if self.deterministic:
|
||||
return torch.Tensor([0.])
|
||||
logtwopi = np.log(2.0 * np.pi)
|
||||
return 0.5 * torch.sum(
|
||||
logtwopi + self.logvar
|
||||
+ torch.pow(sample - self.mean, 2) / self.var,
|
||||
dim=dims)
|
||||
|
||||
def mode(self):
|
||||
return self.mean
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
if self.with_conv:
|
||||
pad = (0, 1, 0, 1)
|
||||
x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
|
||||
x = self.conv(x)
|
||||
else:
|
||||
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
||||
return x
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
in_channels,
|
||||
out_channels=None,
|
||||
conv_shortcut=False,
|
||||
dropout,
|
||||
temb_channels=512):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_conv_shortcut = conv_shortcut
|
||||
|
||||
self.norm1 = Normalize(in_channels)
|
||||
self.conv1 = torch.nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if temb_channels > 0:
|
||||
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
|
||||
self.norm2 = Normalize(out_channels)
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = torch.nn.Conv2d(
|
||||
out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
self.conv_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
else:
|
||||
self.nin_shortcut = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0)
|
||||
|
||||
def forward(self, x, temb):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
if temb is not None:
|
||||
h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
|
||||
|
||||
h = self.norm2(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.dropout(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
if self.use_conv_shortcut:
|
||||
x = self.conv_shortcut(x)
|
||||
else:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = Normalize(in_channels)
|
||||
self.q = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.k = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.v = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.proj_out = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h * w)
|
||||
q = q.permute(0, 2, 1) # b,hw,c
|
||||
k = k.reshape(b, c, h * w) # b,c,hw
|
||||
w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
||||
w_ = w_ * (int(c)**(-0.5))
|
||||
w_ = torch.nn.functional.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h * w)
|
||||
w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
|
||||
h_ = torch.bmm(
|
||||
v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
||||
h_ = h_.reshape(b, c, h, w)
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x + h_
|
||||
|
||||
|
||||
class AttnBlock(nn.Module): # noqa
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = Normalize(in_channels)
|
||||
self.q = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.k = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.v = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.proj_out = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h * w)
|
||||
q = q.permute(0, 2, 1) # b,hw,c
|
||||
k = k.reshape(b, c, h * w) # b,c,hw
|
||||
w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
||||
w_ = w_ * (int(c)**(-0.5))
|
||||
w_ = torch.nn.functional.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h * w)
|
||||
w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
|
||||
h_ = torch.bmm(
|
||||
v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
||||
h_ = h_.reshape(b, c, h, w)
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x + h_
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.nn.functional.interpolate(
|
||||
x, scale_factor=2.0, mode='nearest')
|
||||
if self.with_conv:
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Downsample(nn.Module): # noqa
|
||||
|
||||
def __init__(self, in_channels, with_conv):
|
||||
super().__init__()
|
||||
self.with_conv = with_conv
|
||||
if self.with_conv:
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = torch.nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
if self.with_conv:
|
||||
pad = (0, 1, 0, 1)
|
||||
x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
|
||||
x = self.conv(x)
|
||||
else:
|
||||
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
double_z=True,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignore_kwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
|
||||
# downsampling
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1, ) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in,
|
||||
2 * z_channels if double_z else z_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1], temb)
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
ch,
|
||||
out_ch,
|
||||
ch_mult=(1, 2, 4, 8),
|
||||
num_res_blocks,
|
||||
attn_resolutions,
|
||||
dropout=0.0,
|
||||
resamp_with_conv=True,
|
||||
in_channels,
|
||||
resolution,
|
||||
z_channels,
|
||||
give_pre_end=False,
|
||||
tanh_out=False,
|
||||
use_linear_attn=False,
|
||||
attn_type='vanilla',
|
||||
**ignorekwargs):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.temb_ch = 0
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.give_pre_end = give_pre_end
|
||||
self.tanh_out = tanh_out
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2**(self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
print('Working with z of shape {} = {} dimensions.'.format(
|
||||
self.z_shape, np.prod(self.z_shape)))
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = torch.nn.Conv2d(
|
||||
z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_in,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
block.append(
|
||||
ResnetBlock(
|
||||
in_channels=block_in,
|
||||
out_channels=block_out,
|
||||
temb_channels=self.temb_ch,
|
||||
dropout=dropout))
|
||||
block_in = block_out
|
||||
if curr_res in attn_resolutions:
|
||||
attn.append(AttnBlock(block_in))
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in, resamp_with_conv)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = Normalize(block_in)
|
||||
self.conv_out = torch.nn.Conv2d(
|
||||
block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z):
|
||||
self.last_z_shape = z.shape
|
||||
|
||||
# timestep embedding
|
||||
temb = None
|
||||
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h, temb)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h, temb)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h, temb)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
if self.give_pre_end:
|
||||
return h
|
||||
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
if self.tanh_out:
|
||||
h = torch.tanh(h)
|
||||
return h
|
||||
|
||||
|
||||
class AutoencoderKL(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=[],
|
||||
image_key='image',
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
ema_decay=None,
|
||||
learn_logvar=False):
|
||||
super().__init__()
|
||||
self.learn_logvar = learn_logvar
|
||||
self.image_key = image_key
|
||||
self.encoder = Encoder(**ddconfig)
|
||||
self.decoder = Decoder(**ddconfig)
|
||||
assert ddconfig['double_z']
|
||||
self.quant_conv = torch.nn.Conv2d(2 * ddconfig['z_channels'],
|
||||
2 * embed_dim, 1)
|
||||
self.post_quant_conv = torch.nn.Conv2d(embed_dim,
|
||||
ddconfig['z_channels'], 1)
|
||||
self.embed_dim = embed_dim
|
||||
if colorize_nlabels is not None:
|
||||
assert type(colorize_nlabels) == int
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, colorize_nlabels, 1, 1))
|
||||
if monitor is not None:
|
||||
self.monitor = monitor
|
||||
self.use_ema = ema_decay is not None
|
||||
if ckpt_path is not None:
|
||||
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
||||
|
||||
def init_from_ckpt(self, path, ignore_keys=list()):
|
||||
sd = torch.load(path, map_location='cpu')['state_dict']
|
||||
keys = list(sd.keys())
|
||||
for key in keys:
|
||||
print(key, sd[key].shape)
|
||||
import collections
|
||||
sd_new = collections.OrderedDict()
|
||||
for k in keys:
|
||||
if k.find('first_stage_model') >= 0:
|
||||
k_new = k.split('first_stage_model.')[-1]
|
||||
sd_new[k_new] = sd[k]
|
||||
self.load_state_dict(sd_new, strict=True)
|
||||
print(f'Restored from {path}')
|
||||
|
||||
def init_from_ckpt2(self, path, ignore_keys=list()):
|
||||
sd = torch.load(path, map_location='cpu')['state_dict']
|
||||
keys = list(sd.keys())
|
||||
|
||||
first_stage_model
|
||||
for k in keys:
|
||||
for ik in ignore_keys:
|
||||
if k.startswith(ik):
|
||||
print('Deleting key {} from state_dict.'.format(k))
|
||||
del sd[k]
|
||||
self.load_state_dict(sd, strict=False)
|
||||
print(f'Restored from {path}')
|
||||
|
||||
def on_train_batch_end(self, *args, **kwargs):
|
||||
if self.use_ema:
|
||||
self.model_ema(self)
|
||||
|
||||
def encode(self, x):
|
||||
h = self.encoder(x)
|
||||
moments = self.quant_conv(h)
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
return posterior
|
||||
|
||||
def decode(self, z):
|
||||
z = self.post_quant_conv(z)
|
||||
dec = self.decoder(z)
|
||||
return dec
|
||||
|
||||
def forward(self, input, sample_posterior=True):
|
||||
posterior = self.encode(input)
|
||||
if sample_posterior:
|
||||
z = posterior.sample()
|
||||
else:
|
||||
z = posterior.mode()
|
||||
dec = self.decode(z)
|
||||
return dec, posterior
|
||||
|
||||
def get_input(self, batch, k):
|
||||
x = batch[k]
|
||||
if len(x.shape) == 3:
|
||||
x = x[..., None]
|
||||
x = x.permute(0, 3, 1,
|
||||
2).to(memory_format=torch.contiguous_format).float()
|
||||
return x
|
||||
|
||||
def get_last_layer(self):
|
||||
return self.decoder.conv_out.weight
|
||||
|
||||
@torch.no_grad()
|
||||
def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
if not only_inputs:
|
||||
xrec, posterior = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec.shape[1] > 3
|
||||
x = self.to_rgb(x)
|
||||
xrec = self.to_rgb(xrec)
|
||||
log['samples'] = self.decode(torch.randn_like(posterior.sample()))
|
||||
log['reconstructions'] = xrec
|
||||
if log_ema or self.use_ema:
|
||||
with self.ema_scope():
|
||||
xrec_ema, posterior_ema = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec_ema.shape[1] > 3
|
||||
xrec_ema = self.to_rgb(xrec_ema)
|
||||
log['samples_ema'] = self.decode(
|
||||
torch.randn_like(posterior_ema.sample()))
|
||||
log['reconstructions_ema'] = xrec_ema
|
||||
log['inputs'] = x
|
||||
return log
|
||||
|
||||
def to_rgb(self, x):
|
||||
assert self.image_key == 'segmentation'
|
||||
if not hasattr(self, 'colorize'):
|
||||
self.register_buffer('colorize',
|
||||
torch.randn(3, x.shape[1], 1, 1).to(x))
|
||||
x = F.conv2d(x, weight=self.colorize)
|
||||
x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
|
||||
return x
|
||||
|
||||
|
||||
class IdentityFirstStage(torch.nn.Module):
|
||||
|
||||
def __init__(self, *args, vq_interface=False, **kwargs):
|
||||
self.vq_interface = vq_interface
|
||||
super().__init__()
|
||||
|
||||
def encode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def decode(self, x, *args, **kwargs):
|
||||
return x
|
||||
|
||||
def quantize(self, x, *args, **kwargs):
|
||||
if self.vq_interface:
|
||||
return x, None, [None, None, None]
|
||||
return x
|
||||
|
||||
def forward(self, x, *args, **kwargs):
|
||||
return x
|
||||
143
modelscope/models/multi_modal/videocomposer/clip.py
Normal file
143
modelscope/models/multi_modal/videocomposer/clip.py
Normal file
@@ -0,0 +1,143 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import numpy as np
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as T
|
||||
|
||||
|
||||
class FrozenOpenCLIPEmbedder(nn.Module):
|
||||
"""
|
||||
Uses the OpenCLIP transformer encoder for text
|
||||
"""
|
||||
LAYERS = ['last', 'penultimate']
|
||||
|
||||
def __init__(self,
|
||||
arch='ViT-H-14',
|
||||
pretrained='laion2b_s32b_b79k',
|
||||
device='cuda',
|
||||
max_length=77,
|
||||
freeze=True,
|
||||
layer='last'):
|
||||
super().__init__()
|
||||
assert layer in self.LAYERS
|
||||
model, _, _ = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=pretrained)
|
||||
del model.visual
|
||||
self.model = model
|
||||
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
if freeze:
|
||||
self.freeze()
|
||||
self.layer = layer
|
||||
if self.layer == 'last':
|
||||
self.layer_idx = 0
|
||||
elif self.layer == 'penultimate':
|
||||
self.layer_idx = 1
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
def freeze(self):
|
||||
self.model = self.model.eval()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, text):
|
||||
tokens = open_clip.tokenize(text)
|
||||
z = self.encode_with_transformer(tokens.to(self.device))
|
||||
return z
|
||||
|
||||
def encode_with_transformer(self, text):
|
||||
x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model]
|
||||
x = x + self.model.positional_embedding
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.model.ln_final(x)
|
||||
return x
|
||||
|
||||
def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
|
||||
for i, r in enumerate(self.model.transformer.resblocks):
|
||||
if i == len(self.model.transformer.resblocks) - self.layer_idx:
|
||||
break
|
||||
if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(
|
||||
):
|
||||
x = checkpoint(r, x, attn_mask)
|
||||
else:
|
||||
x = r(x, attn_mask=attn_mask)
|
||||
return x
|
||||
|
||||
def encode(self, text):
|
||||
return self(text)
|
||||
|
||||
|
||||
class FrozenOpenCLIPVisualEmbedder(nn.Module):
|
||||
"""
|
||||
Uses the OpenCLIP transformer encoder for text
|
||||
"""
|
||||
LAYERS = ['last', 'penultimate']
|
||||
|
||||
def __init__(self,
|
||||
arch='ViT-H-14',
|
||||
pretrained='laion2b_s32b_b79k',
|
||||
device='cuda',
|
||||
max_length=77,
|
||||
freeze=True,
|
||||
layer='last',
|
||||
input_shape=(224, 224, 3)):
|
||||
super().__init__()
|
||||
assert layer in self.LAYERS
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
arch, device=torch.device('cpu'), pretrained=pretrained)
|
||||
del model.transformer
|
||||
self.model = model
|
||||
data_white = np.ones(input_shape, dtype=np.uint8) * 255
|
||||
self.black_image = preprocess(T.ToPILImage()(data_white)).unsqueeze(0)
|
||||
self.preprocess = preprocess
|
||||
|
||||
self.device = device
|
||||
self.max_length = max_length # 77
|
||||
if freeze:
|
||||
self.freeze()
|
||||
self.layer = layer # 'penultimate'
|
||||
if self.layer == 'last':
|
||||
self.layer_idx = 0
|
||||
elif self.layer == 'penultimate':
|
||||
self.layer_idx = 1
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
def freeze(self):
|
||||
self.model = self.model.eval()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, image):
|
||||
# tokens = open_clip.tokenize(text)
|
||||
z = self.model.encode_image(image.to(self.device))
|
||||
return z
|
||||
|
||||
def encode_with_transformer(self, text):
|
||||
x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model]
|
||||
x = x + self.model.positional_embedding
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.model.ln_final(x)
|
||||
return x
|
||||
|
||||
def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
|
||||
for i, r in enumerate(self.model.transformer.resblocks):
|
||||
if i == len(self.model.transformer.resblocks) - self.layer_idx:
|
||||
break
|
||||
if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(
|
||||
):
|
||||
x = checkpoint(r, x, attn_mask)
|
||||
else:
|
||||
x = r(x, attn_mask=attn_mask)
|
||||
return x
|
||||
|
||||
def encode(self, text):
|
||||
return self(text)
|
||||
156
modelscope/models/multi_modal/videocomposer/config.py
Normal file
156
modelscope/models/multi_modal/videocomposer/config.py
Normal file
@@ -0,0 +1,156 @@
|
||||
import logging
|
||||
import os
|
||||
import os.path as osp
|
||||
from datetime import datetime
|
||||
|
||||
import torch
|
||||
from easydict import EasyDict
|
||||
|
||||
cfg = EasyDict(__name__='Config: VideoComposer')
|
||||
|
||||
pmi_world_size = int(os.getenv('WORLD_SIZE', 1))
|
||||
gpus_per_machine = torch.cuda.device_count()
|
||||
world_size = pmi_world_size * gpus_per_machine
|
||||
|
||||
cfg.video_compositions = [
|
||||
'text', 'mask', 'depthmap', 'sketch', 'motion', 'image', 'local_image',
|
||||
'single_sketch'
|
||||
]
|
||||
|
||||
# dataset
|
||||
cfg.root_dir = 'webvid10m/'
|
||||
|
||||
cfg.alpha = 0.7
|
||||
|
||||
cfg.misc_size = 384
|
||||
cfg.depth_std = 20.0
|
||||
cfg.depth_clamp = 10.0
|
||||
cfg.hist_sigma = 10.0
|
||||
|
||||
cfg.use_image_dataset = False
|
||||
cfg.alpha_img = 0.7
|
||||
|
||||
cfg.resolution = 256
|
||||
cfg.mean = [0.5, 0.5, 0.5]
|
||||
cfg.std = [0.5, 0.5, 0.5]
|
||||
|
||||
# sketch
|
||||
cfg.sketch_mean = [0.485, 0.456, 0.406]
|
||||
cfg.sketch_std = [0.229, 0.224, 0.225]
|
||||
|
||||
# dataloader
|
||||
cfg.max_words = 1000
|
||||
|
||||
cfg.frame_lens = [
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
]
|
||||
cfg.feature_framerates = [
|
||||
4,
|
||||
]
|
||||
cfg.feature_framerate = 4
|
||||
cfg.batch_sizes = {
|
||||
str(1): 1,
|
||||
str(4): 1,
|
||||
str(8): 1,
|
||||
str(16): 1,
|
||||
}
|
||||
|
||||
cfg.chunk_size = 64
|
||||
cfg.num_workers = 8
|
||||
cfg.prefetch_factor = 2
|
||||
cfg.seed = 8888
|
||||
|
||||
# diffusion
|
||||
cfg.num_timesteps = 1000
|
||||
cfg.mean_type = 'eps'
|
||||
cfg.var_type = 'fixed_small'
|
||||
cfg.loss_type = 'mse'
|
||||
cfg.ddim_timesteps = 50
|
||||
cfg.ddim_eta = 0.0
|
||||
cfg.clamp = 1.0
|
||||
cfg.share_noise = False
|
||||
cfg.use_div_loss = False
|
||||
|
||||
# classifier-free guidance
|
||||
cfg.p_zero = 0.9
|
||||
cfg.guide_scale = 6.0
|
||||
|
||||
# stabel diffusion
|
||||
cfg.sd_checkpoint = 'v2-1_512-ema-pruned.ckpt'
|
||||
|
||||
# clip vision encoder
|
||||
cfg.vit_image_size = 336
|
||||
cfg.vit_patch_size = 14
|
||||
cfg.vit_dim = 1024
|
||||
cfg.vit_out_dim = 768
|
||||
cfg.vit_heads = 16
|
||||
cfg.vit_layers = 24
|
||||
cfg.vit_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||
cfg.vit_std = [0.26862954, 0.26130258, 0.27577711]
|
||||
cfg.clip_checkpoint = 'open_clip_pytorch_model.bin'
|
||||
cfg.mvs_visual = False
|
||||
|
||||
# unet
|
||||
cfg.unet_in_dim = 4
|
||||
cfg.unet_concat_dim = 8
|
||||
cfg.unet_y_dim = cfg.vit_out_dim
|
||||
cfg.unet_context_dim = 1024
|
||||
cfg.unet_out_dim = 8 if cfg.var_type.startswith('learned') else 4
|
||||
cfg.unet_dim = 320
|
||||
cfg.unet_dim_mult = [1, 2, 4, 4]
|
||||
cfg.unet_res_blocks = 2
|
||||
cfg.unet_num_heads = 8
|
||||
cfg.unet_head_dim = 64
|
||||
cfg.unet_attn_scales = [1 / 1, 1 / 2, 1 / 4]
|
||||
cfg.unet_dropout = 0.1
|
||||
cfg.misc_dropout = 0.5
|
||||
cfg.p_all_zero = 0.1
|
||||
cfg.p_all_keep = 0.1
|
||||
cfg.temporal_conv = False
|
||||
cfg.temporal_attn_times = 1
|
||||
cfg.temporal_attention = True
|
||||
|
||||
cfg.use_fps_condition = False
|
||||
cfg.use_sim_mask = False
|
||||
|
||||
# Default: load 2d pretrain
|
||||
cfg.pretrained = False
|
||||
cfg.fix_weight = False
|
||||
|
||||
# Default resume
|
||||
cfg.resume = True
|
||||
cfg.resume_step = 148000
|
||||
cfg.resume_check_dir = '.'
|
||||
cfg.resume_checkpoint = os.path.join(
|
||||
cfg.resume_check_dir,
|
||||
f'step_{cfg.resume_step}/non_ema_{cfg.resume_step}.pth')
|
||||
cfg.resume_optimizer = False
|
||||
if cfg.resume_optimizer:
|
||||
cfg.resume_optimizer = os.path.join(
|
||||
cfg.resume_check_dir, f'optimizer_step_{cfg.resume_step}.pt')
|
||||
|
||||
# acceleration
|
||||
cfg.use_ema = True
|
||||
# for debug, no ema
|
||||
if world_size < 2:
|
||||
cfg.use_ema = False
|
||||
cfg.load_from = None
|
||||
|
||||
cfg.use_checkpoint = True
|
||||
cfg.use_sharded_ddp = False
|
||||
cfg.use_fsdp = False
|
||||
cfg.use_fp16 = True
|
||||
|
||||
# training
|
||||
cfg.ema_decay = 0.9999
|
||||
cfg.viz_interval = 1000
|
||||
cfg.save_ckp_interval = 1000
|
||||
|
||||
# logging
|
||||
cfg.log_interval = 100
|
||||
composition_strings = '_'.join(cfg.video_compositions)
|
||||
# Default log_dir
|
||||
cfg.log_dir = 'outputs/'
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .samplers import *
|
||||
from .tokenizers import *
|
||||
from .transforms import *
|
||||
158
modelscope/models/multi_modal/videocomposer/data/samplers.py
Normal file
158
modelscope/models/multi_modal/videocomposer/data/samplers.py
Normal file
@@ -0,0 +1,158 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import os.path as osp
|
||||
|
||||
import json
|
||||
import numpy as np
|
||||
from torch.utils.data.sampler import Sampler
|
||||
|
||||
from modelscope.models.multi_modal.videocomposer.ops.distributed import (
|
||||
get_rank, get_world_size, shared_random_seed)
|
||||
from modelscope.models.multi_modal.videocomposer.ops.utils import (ceil_divide,
|
||||
read)
|
||||
|
||||
__all__ = ['BatchSampler', 'GroupSampler', 'ImgGroupSampler']
|
||||
|
||||
|
||||
class BatchSampler(Sampler):
|
||||
r"""An infinite batch sampler.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
dataset_size,
|
||||
batch_size,
|
||||
num_replicas=None,
|
||||
rank=None,
|
||||
shuffle=False,
|
||||
seed=None):
|
||||
self.dataset_size = dataset_size
|
||||
self.batch_size = batch_size
|
||||
self.num_replicas = num_replicas or get_world_size()
|
||||
self.rank = rank or get_rank()
|
||||
self.shuffle = shuffle
|
||||
self.seed = seed or shared_random_seed()
|
||||
self.rng = np.random.default_rng(self.seed + self.rank)
|
||||
self.batches_per_rank = ceil_divide(
|
||||
dataset_size, self.num_replicas * self.batch_size)
|
||||
self.samples_per_rank = self.batches_per_rank * self.batch_size
|
||||
|
||||
# rank indices
|
||||
indices = self.rng.permutation(
|
||||
self.samples_per_rank) if shuffle else np.arange(
|
||||
self.samples_per_rank)
|
||||
indices = indices * self.num_replicas + self.rank
|
||||
indices = indices[indices < dataset_size]
|
||||
self.indices = indices
|
||||
|
||||
def __iter__(self):
|
||||
start = 0
|
||||
while True:
|
||||
batch = [
|
||||
self.indices[i % len(self.indices)]
|
||||
for i in range(start, start + self.batch_size)
|
||||
]
|
||||
if self.shuffle and (start + self.batch_size) > len(self.indices):
|
||||
self.rng.shuffle(self.indices)
|
||||
start = (start + self.batch_size) % len(self.indices)
|
||||
yield batch
|
||||
|
||||
|
||||
class GroupSampler(Sampler):
|
||||
|
||||
def __init__(self,
|
||||
group_file,
|
||||
batch_size,
|
||||
alpha=0.7,
|
||||
update_interval=5000,
|
||||
seed=8888):
|
||||
self.group_file = group_file
|
||||
self.group_folder = osp.join(osp.dirname(group_file), 'groups')
|
||||
self.batch_size = batch_size
|
||||
self.alpha = alpha
|
||||
self.update_interval = update_interval
|
||||
self.seed = seed
|
||||
self.rng = np.random.default_rng(seed)
|
||||
|
||||
def __iter__(self):
|
||||
while True:
|
||||
# keep groups up-to-date
|
||||
self.update_groups()
|
||||
|
||||
# collect items
|
||||
items = self.sample()
|
||||
while len(items) < self.batch_size:
|
||||
items += self.sample()
|
||||
|
||||
# sample a batch
|
||||
batch = self.rng.choice(
|
||||
items,
|
||||
self.batch_size,
|
||||
replace=False if len(items) >= self.batch_size else True)
|
||||
yield [u.strip().split(',') for u in batch]
|
||||
|
||||
def update_groups(self):
|
||||
if not hasattr(self, '_step'):
|
||||
self._step = 0
|
||||
if self._step % self.update_interval == 0:
|
||||
self.groups = json.loads(read(self.group_file))
|
||||
self._step += 1
|
||||
|
||||
def sample(self):
|
||||
scales = np.array(
|
||||
[float(next(iter(u)).split(':')[-1]) for u in self.groups])
|
||||
p = scales**self.alpha / (scales**self.alpha).sum()
|
||||
group = self.rng.choice(self.groups, p=p)
|
||||
list_file = osp.join(self.group_folder,
|
||||
self.rng.choice(next(iter(group.values()))))
|
||||
return read(list_file).strip().split('\n')
|
||||
|
||||
|
||||
class ImgGroupSampler(Sampler):
|
||||
|
||||
def __init__(self,
|
||||
group_file,
|
||||
batch_size,
|
||||
alpha=0.7,
|
||||
update_interval=5000,
|
||||
seed=8888):
|
||||
self.group_file = group_file
|
||||
self.group_folder = osp.join(osp.dirname(group_file), 'groups')
|
||||
self.batch_size = batch_size
|
||||
self.alpha = alpha
|
||||
self.update_interval = update_interval
|
||||
self.seed = seed
|
||||
self.rng = np.random.default_rng(seed)
|
||||
|
||||
def __iter__(self):
|
||||
while True:
|
||||
# keep groups up-to-date
|
||||
self.update_groups()
|
||||
|
||||
# collect items
|
||||
items = self.sample()
|
||||
while len(items) < self.batch_size:
|
||||
items += self.sample()
|
||||
|
||||
# sample a batch
|
||||
batch = self.rng.choice(
|
||||
items,
|
||||
self.batch_size,
|
||||
replace=False if len(items) >= self.batch_size else True)
|
||||
yield [u.strip().split(',', 1) for u in batch]
|
||||
|
||||
def update_groups(self):
|
||||
if not hasattr(self, '_step'):
|
||||
self._step = 0
|
||||
if self._step % self.update_interval == 0:
|
||||
self.groups = json.loads(read(self.group_file))
|
||||
|
||||
self._step += 1
|
||||
|
||||
def sample(self):
|
||||
scales = np.array(
|
||||
[float(next(iter(u)).split(':')[-1]) for u in self.groups])
|
||||
p = scales**self.alpha / (scales**self.alpha).sum()
|
||||
group = self.rng.choice(self.groups, p=p)
|
||||
list_file = osp.join(self.group_folder,
|
||||
self.rng.choice(next(iter(group.values()))))
|
||||
return read(list_file).strip().split('\n')
|
||||
184
modelscope/models/multi_modal/videocomposer/data/tokenizers.py
Normal file
184
modelscope/models/multi_modal/videocomposer/data/tokenizers.py
Normal file
@@ -0,0 +1,184 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import gzip
|
||||
import html
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
import torch
|
||||
from tokenizers import BertWordPieceTokenizer, CharBPETokenizer
|
||||
|
||||
__all__ = ['CLIPTokenizer']
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def default_bpe():
|
||||
root = os.path.realpath(__file__)
|
||||
root = '/'.join(root.split('/')[:-1])
|
||||
return os.path.join(root, 'bpe_simple_vocab_16e6.txt.gz')
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def bytes_to_unicode():
|
||||
"""
|
||||
Returns list of utf-8 byte and a corresponding list of unicode strings.
|
||||
The reversible bpe codes work on unicode strings.
|
||||
This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
|
||||
When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
|
||||
This is a signficant percentage of your normal, say, 32K bpe vocab.
|
||||
To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
|
||||
And avoids mapping to whitespace/control characters the bpe code barfs on.
|
||||
"""
|
||||
bs = list(range(ord('!'),
|
||||
ord('~') + 1)) + list(range(
|
||||
ord('¡'),
|
||||
ord('¬') + 1)) + list(range(ord('®'),
|
||||
ord('ÿ') + 1))
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8 + n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
|
||||
def get_pairs(word):
|
||||
"""Return set of symbol pairs in a word.
|
||||
Word is represented as tuple of symbols (symbols being variable-length strings).
|
||||
"""
|
||||
pairs = set()
|
||||
prev_char = word[0]
|
||||
for char in word[1:]:
|
||||
pairs.add((prev_char, char))
|
||||
prev_char = char
|
||||
return pairs
|
||||
|
||||
|
||||
def basic_clean(text):
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def whitespace_clean(text):
|
||||
text = re.sub(r'\s+', ' ', text)
|
||||
text = text.strip()
|
||||
return text
|
||||
|
||||
|
||||
class SimpleTokenizer(object):
|
||||
|
||||
def __init__(self, bpe_path: str = default_bpe()):
|
||||
self.byte_encoder = bytes_to_unicode()
|
||||
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
||||
merges = gzip.open(bpe_path).read().decode('utf-8').split('\n')
|
||||
merges = merges[1:49152 - 256 - 2 + 1]
|
||||
merges = [tuple(merge.split()) for merge in merges]
|
||||
vocab = list(bytes_to_unicode().values())
|
||||
vocab = vocab + [v + '</w>' for v in vocab]
|
||||
for merge in merges:
|
||||
vocab.append(''.join(merge))
|
||||
vocab.extend(['<|startoftext|>', '<|endoftext|>'])
|
||||
self.encoder = dict(zip(vocab, range(len(vocab))))
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
||||
self.cache = {
|
||||
'<|startoftext|>': '<|startoftext|>',
|
||||
'<|endoftext|>': '<|endoftext|>'
|
||||
}
|
||||
self.pat = re.compile(
|
||||
r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
|
||||
re.IGNORECASE)
|
||||
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
return self.cache[token]
|
||||
word = tuple(token[:-1]) + (token[-1] + '</w>', )
|
||||
pairs = get_pairs(word)
|
||||
|
||||
if not pairs:
|
||||
return token + '</w>'
|
||||
|
||||
while True:
|
||||
bigram = min(
|
||||
pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf')))
|
||||
if bigram not in self.bpe_ranks:
|
||||
break
|
||||
first, second = bigram
|
||||
new_word = []
|
||||
i = 0
|
||||
while i < len(word):
|
||||
try:
|
||||
j = word.index(first, i)
|
||||
new_word.extend(word[i:j])
|
||||
i = j
|
||||
except Exception as e:
|
||||
new_word.extend(word[i:])
|
||||
print(e)
|
||||
break
|
||||
|
||||
if word[i] == first and i < len(word) - 1 and word[
|
||||
i + 1] == second:
|
||||
new_word.append(first + second)
|
||||
i += 2
|
||||
else:
|
||||
new_word.append(word[i])
|
||||
i += 1
|
||||
new_word = tuple(new_word)
|
||||
word = new_word
|
||||
if len(word) == 1:
|
||||
break
|
||||
else:
|
||||
pairs = get_pairs(word)
|
||||
word = ' '.join(word)
|
||||
self.cache[token] = word
|
||||
return word
|
||||
|
||||
def encode(self, text):
|
||||
bpe_tokens = []
|
||||
text = whitespace_clean(basic_clean(text)).lower()
|
||||
for token in re.findall(self.pat, text):
|
||||
token = ''.join(self.byte_encoder[b]
|
||||
for b in token.encode('utf-8'))
|
||||
bpe_tokens.extend(self.encoder[bpe_token]
|
||||
for bpe_token in self.bpe(token).split(' '))
|
||||
return bpe_tokens
|
||||
|
||||
def decode(self, tokens):
|
||||
text = ''.join([self.decoder[token] for token in tokens])
|
||||
text = bytearray([self.byte_decoder[c] for c in text]).decode(
|
||||
'utf-8', errors='replace').replace('</w>', ' ')
|
||||
return text
|
||||
|
||||
|
||||
class CLIPTokenizer(object):
|
||||
|
||||
def __init__(self, length=77):
|
||||
self.length = length
|
||||
|
||||
# init tokenizer
|
||||
self.tokenizer = SimpleTokenizer(bpe_path=default_bpe())
|
||||
self.sos_token = self.tokenizer.encoder['<|startoftext|>']
|
||||
self.eos_token = self.tokenizer.encoder['<|endoftext|>']
|
||||
self.vocab_size = len(self.tokenizer.encoder)
|
||||
|
||||
def __call__(self, sequence):
|
||||
if isinstance(sequence, str):
|
||||
return torch.LongTensor(self._tokenizer(sequence))
|
||||
elif isinstance(sequence, list):
|
||||
return torch.LongTensor([self._tokenizer(u) for u in sequence])
|
||||
else:
|
||||
raise TypeError(
|
||||
f'Expected the "sequence" to be a string or a list, but got {type(sequence)}'
|
||||
)
|
||||
|
||||
def _tokenizer(self, text):
|
||||
tokens = self.tokenizer.encode(text)[:self.length - 2]
|
||||
tokens = [self.sos_token] + tokens + [self.eos_token]
|
||||
tokens = tokens + [0] * (self.length - len(tokens))
|
||||
return tokens
|
||||
400
modelscope/models/multi_modal/videocomposer/data/transforms.py
Normal file
400
modelscope/models/multi_modal/videocomposer/data/transforms.py
Normal file
@@ -0,0 +1,400 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
from PIL import Image, ImageFilter
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
__all__ = [
|
||||
'Compose', 'Resize', 'Rescale', 'CenterCrop', 'CenterCropV2', 'RandomCrop',
|
||||
'RandomCropV2', 'RandomHFlip', 'GaussianBlur', 'ColorJitter', 'RandomGray',
|
||||
'ToTensor', 'Normalize', 'ResizeRandomCrop', 'ExtractResizeRandomCrop',
|
||||
'ExtractResizeAssignCrop'
|
||||
]
|
||||
|
||||
|
||||
def random_resize(img, size):
|
||||
img = [
|
||||
TF.resize(
|
||||
u,
|
||||
size,
|
||||
interpolation=random.choice([
|
||||
InterpolationMode.BILINEAR, InterpolationMode.BICUBIC,
|
||||
InterpolationMode.LANCZOS
|
||||
])) for u in img
|
||||
]
|
||||
return img
|
||||
|
||||
|
||||
class CenterCropV3(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
# fast resize
|
||||
while min(img.size) >= 2 * self.size:
|
||||
img = img.resize((img.width // 2, img.height // 2),
|
||||
resample=Image.BOX)
|
||||
scale = self.size / min(img.size)
|
||||
img = img.resize((round(scale * img.width), round(scale * img.height)),
|
||||
resample=Image.BICUBIC)
|
||||
|
||||
# center crop
|
||||
x1 = (img.width - self.size) // 2
|
||||
y1 = (img.height - self.size) // 2
|
||||
img = img.crop((x1, y1, x1 + self.size, y1 + self.size))
|
||||
return img
|
||||
|
||||
|
||||
class Compose(object):
|
||||
|
||||
def __init__(self, transforms):
|
||||
self.transforms = transforms
|
||||
|
||||
def __getitem__(self, index):
|
||||
if isinstance(index, slice):
|
||||
return Compose(self.transforms[index])
|
||||
else:
|
||||
return self.transforms[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.transforms)
|
||||
|
||||
def __call__(self, rgb):
|
||||
for t in self.transforms:
|
||||
rgb = t(rgb)
|
||||
return rgb
|
||||
|
||||
|
||||
class Resize(object):
|
||||
|
||||
def __init__(self, size=256):
|
||||
if isinstance(size, int):
|
||||
size = (size, size)
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
rgb = [u.resize(self.size, Image.BILINEAR) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class Rescale(object):
|
||||
|
||||
def __init__(self, size=256, interpolation=Image.BILINEAR):
|
||||
self.size = size
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
scale = self.size / min(w, h)
|
||||
out_w, out_h = int(round(w * scale)), int(round(h * scale))
|
||||
rgb = [u.resize((out_w, out_h), self.interpolation) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCrop(object):
|
||||
|
||||
def __init__(self, size=224):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, rgb):
|
||||
w, h = rgb[0].size
|
||||
assert min(w, h) >= self.size
|
||||
x1 = (w - self.size) // 2
|
||||
y1 = (h - self.size) // 2
|
||||
rgb = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
# self.min_area = min_area
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size # (518, 292)
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class ExtractResizeRandomCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
out_w = self.size
|
||||
out_h = self.size
|
||||
w, h = rgb[0].size # (518, 292)
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
wh = [x1, y1, x1 + out_w, y1 + out_h]
|
||||
|
||||
return rgb, wh
|
||||
|
||||
|
||||
class ExtractResizeAssignCrop(object):
|
||||
|
||||
def __init__(self, size=256, size_short=292):
|
||||
self.size = size
|
||||
self.size_short = size_short
|
||||
|
||||
def __call__(self, rgb, wh):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
while min(rgb[0].size) >= 2 * self.size_short:
|
||||
rgb = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in rgb
|
||||
]
|
||||
scale = self.size_short / min(rgb[0].size)
|
||||
rgb = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in rgb
|
||||
]
|
||||
|
||||
rgb = [u.crop(wh) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class CenterCropV2(object):
|
||||
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img):
|
||||
# fast resize
|
||||
while min(img[0].size) >= 2 * self.size:
|
||||
img = [
|
||||
u.resize((u.width // 2, u.height // 2), resample=Image.BOX)
|
||||
for u in img
|
||||
]
|
||||
scale = self.size / min(img[0].size)
|
||||
img = [
|
||||
u.resize((round(scale * u.width), round(scale * u.height)),
|
||||
resample=Image.BICUBIC) for u in img
|
||||
]
|
||||
|
||||
# center crop
|
||||
x1 = (img[0].width - self.size) // 2
|
||||
y1 = (img[0].height - self.size) // 2
|
||||
img = [u.crop((x1, y1, x1 + self.size, y1 + self.size)) for u in img]
|
||||
return img
|
||||
|
||||
|
||||
class RandomCrop(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4):
|
||||
self.size = size
|
||||
self.min_area = min_area
|
||||
|
||||
def __call__(self, rgb):
|
||||
|
||||
# consistent crop between rgb and m
|
||||
w, h = rgb[0].size
|
||||
area = w * h
|
||||
out_w, out_h = float('inf'), float('inf')
|
||||
while out_w > w or out_h > h:
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
aspect_ratio = random.uniform(3. / 4., 4. / 3.)
|
||||
out_w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
out_h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
x1 = random.randint(0, w - out_w)
|
||||
y1 = random.randint(0, h - out_h)
|
||||
|
||||
rgb = [u.crop((x1, y1, x1 + out_w, y1 + out_h)) for u in rgb]
|
||||
rgb = [u.resize((self.size, self.size), Image.BILINEAR) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomCropV2(object):
|
||||
|
||||
def __init__(self, size=224, min_area=0.4, ratio=(3. / 4., 4. / 3.)):
|
||||
if isinstance(size, (tuple, list)):
|
||||
self.size = size
|
||||
else:
|
||||
self.size = (size, size)
|
||||
self.min_area = min_area
|
||||
self.ratio = ratio
|
||||
|
||||
def _get_params(self, img):
|
||||
width, height = img.size
|
||||
area = height * width
|
||||
|
||||
for _ in range(10):
|
||||
target_area = random.uniform(self.min_area, 1.0) * area
|
||||
log_ratio = (math.log(self.ratio[0]), math.log(self.ratio[1]))
|
||||
aspect_ratio = math.exp(random.uniform(*log_ratio))
|
||||
|
||||
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
|
||||
if 0 < w <= width and 0 < h <= height:
|
||||
i = random.randint(0, height - h)
|
||||
j = random.randint(0, width - w)
|
||||
return i, j, h, w
|
||||
|
||||
# Fallback to central crop
|
||||
in_ratio = float(width) / float(height)
|
||||
if (in_ratio < min(self.ratio)):
|
||||
w = width
|
||||
h = int(round(w / min(self.ratio)))
|
||||
elif (in_ratio > max(self.ratio)):
|
||||
h = height
|
||||
w = int(round(h * max(self.ratio)))
|
||||
else: # whole image
|
||||
w = width
|
||||
h = height
|
||||
i = (height - h) // 2
|
||||
j = (width - w) // 2
|
||||
return i, j, h, w
|
||||
|
||||
def __call__(self, rgb):
|
||||
i, j, h, w = self._get_params(rgb[0])
|
||||
rgb = [F.resized_crop(u, i, j, h, w, self.size) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class RandomHFlip(object):
|
||||
|
||||
def __init__(self, p=0.5):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.transpose(Image.FLIP_LEFT_RIGHT) for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class GaussianBlur(object):
|
||||
|
||||
def __init__(self, sigmas=[0.1, 2.0], p=0.5):
|
||||
self.sigmas = sigmas
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
sigma = random.uniform(*self.sigmas)
|
||||
rgb = [
|
||||
u.filter(ImageFilter.GaussianBlur(radius=sigma)) for u in rgb
|
||||
]
|
||||
return rgb
|
||||
|
||||
|
||||
class ColorJitter(object):
|
||||
|
||||
def __init__(self,
|
||||
brightness=0.4,
|
||||
contrast=0.4,
|
||||
saturation=0.4,
|
||||
hue=0.1,
|
||||
p=0.5):
|
||||
self.brightness = brightness
|
||||
self.contrast = contrast
|
||||
self.saturation = saturation
|
||||
self.hue = hue
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
brightness, contrast, saturation, hue = self._random_params()
|
||||
transforms = [
|
||||
lambda f: F.adjust_brightness(f, brightness),
|
||||
lambda f: F.adjust_contrast(f, contrast),
|
||||
lambda f: F.adjust_saturation(f, saturation),
|
||||
lambda f: F.adjust_hue(f, hue)
|
||||
]
|
||||
random.shuffle(transforms)
|
||||
for t in transforms:
|
||||
rgb = [t(u) for u in rgb]
|
||||
|
||||
return rgb
|
||||
|
||||
def _random_params(self):
|
||||
brightness = random.uniform(
|
||||
max(0, 1 - self.brightness), 1 + self.brightness)
|
||||
contrast = random.uniform(max(0, 1 - self.contrast), 1 + self.contrast)
|
||||
saturation = random.uniform(
|
||||
max(0, 1 - self.saturation), 1 + self.saturation)
|
||||
hue = random.uniform(-self.hue, self.hue)
|
||||
return brightness, contrast, saturation, hue
|
||||
|
||||
|
||||
class RandomGray(object):
|
||||
|
||||
def __init__(self, p=0.2):
|
||||
self.p = p
|
||||
|
||||
def __call__(self, rgb):
|
||||
if random.random() < self.p:
|
||||
rgb = [u.convert('L').convert('RGB') for u in rgb]
|
||||
return rgb
|
||||
|
||||
|
||||
class ToTensor(object):
|
||||
|
||||
def __call__(self, rgb):
|
||||
rgb = torch.stack([F.to_tensor(u) for u in rgb], dim=0)
|
||||
return rgb
|
||||
|
||||
|
||||
class Normalize(object):
|
||||
|
||||
def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
def __call__(self, rgb):
|
||||
rgb = rgb.clone()
|
||||
rgb.clamp_(0, 1)
|
||||
if not isinstance(self.mean, torch.Tensor):
|
||||
self.mean = rgb.new_tensor(self.mean).view(-1)
|
||||
if not isinstance(self.std, torch.Tensor):
|
||||
self.std = rgb.new_tensor(self.std).view(-1)
|
||||
rgb.sub_(self.mean.view(1, -1, 1, 1)).div_(self.std.view(1, -1, 1, 1))
|
||||
return rgb
|
||||
1514
modelscope/models/multi_modal/videocomposer/diffusion.py
Normal file
1514
modelscope/models/multi_modal/videocomposer/diffusion.py
Normal file
File diff suppressed because it is too large
Load Diff
1697
modelscope/models/multi_modal/videocomposer/dpm_solver.py
Normal file
1697
modelscope/models/multi_modal/videocomposer/dpm_solver.py
Normal file
File diff suppressed because it is too large
Load Diff
120
modelscope/models/multi_modal/videocomposer/mha_flash.py
Normal file
120
modelscope/models/multi_modal/videocomposer/mha_flash.py
Normal file
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from flash_attn.flash_attention import FlashAttention
|
||||
|
||||
|
||||
class FlashAttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
context_dim=None,
|
||||
num_heads=None,
|
||||
head_dim=None,
|
||||
batch_size=4):
|
||||
# consider head_dim first, then num_heads
|
||||
num_heads = dim // head_dim if head_dim else num_heads
|
||||
head_dim = dim // num_heads
|
||||
assert num_heads * head_dim == dim
|
||||
super(FlashAttentionBlock, self).__init__()
|
||||
self.dim = dim
|
||||
self.context_dim = context_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.scale = math.pow(head_dim, -0.25)
|
||||
|
||||
# layers
|
||||
self.norm = nn.GroupNorm(32, dim)
|
||||
self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
|
||||
if context_dim is not None:
|
||||
self.context_kv = nn.Linear(context_dim, dim * 2)
|
||||
self.proj = nn.Conv2d(dim, dim, 1)
|
||||
|
||||
if self.head_dim <= 128 and (self.head_dim % 8) == 0:
|
||||
self.flash_attn = FlashAttention(
|
||||
softmax_scale=None, attention_dropout=0.0)
|
||||
|
||||
# zero out the last layer params
|
||||
nn.init.zeros_(self.proj.weight)
|
||||
|
||||
def _init_weight(self, module):
|
||||
if isinstance(module, nn.Linear):
|
||||
module.weight.data.normal_(mean=0.0, std=0.15)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.Conv2d):
|
||||
module.weight.data.normal_(mean=0.0, std=0.15)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
|
||||
def forward(self, x, context=None):
|
||||
r"""x: [B, C, H, W].
|
||||
context: [B, L, C] or None.
|
||||
"""
|
||||
identity = x
|
||||
b, c, h, w, n, d = *x.size(), self.num_heads, self.head_dim
|
||||
|
||||
# compute query, key, value
|
||||
x = self.norm(x)
|
||||
q, k, v = self.to_qkv(x).view(b, n * 3, d, h * w).chunk(3, dim=1)
|
||||
if context is not None:
|
||||
ck, cv = self.context_kv(context).reshape(b, -1, n * 2,
|
||||
d).permute(0, 2, 3,
|
||||
1).chunk(
|
||||
2, dim=1)
|
||||
k = torch.cat([ck, k], dim=-1)
|
||||
v = torch.cat([cv, v], dim=-1)
|
||||
cq = torch.zeros([b, n, d, 4], dtype=q.dtype, device=q.device)
|
||||
q = torch.cat([q, cq], dim=-1)
|
||||
|
||||
qkv = torch.cat([q, k, v], dim=1)
|
||||
origin_dtype = qkv.dtype
|
||||
qkv = qkv.permute(0, 3, 1, 2).reshape(b, -1, 3, n,
|
||||
d).half().contiguous()
|
||||
out, _ = self.flash_attn(qkv)
|
||||
out.to(origin_dtype)
|
||||
|
||||
if context is not None:
|
||||
out = out[:, :-4, :, :]
|
||||
out = out.permute(0, 2, 3, 1).reshape(b, c, h, w)
|
||||
|
||||
# output
|
||||
x = self.proj(out)
|
||||
return x + identity
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
batch_size = 8
|
||||
flash_net = FlashAttentionBlock(
|
||||
dim=1280,
|
||||
context_dim=512,
|
||||
num_heads=None,
|
||||
head_dim=64,
|
||||
batch_size=batch_size).cuda()
|
||||
|
||||
x = torch.randn([batch_size, 1280, 32, 32], dtype=torch.float32).cuda()
|
||||
context = torch.randn([batch_size, 4, 512], dtype=torch.float32).cuda()
|
||||
# context = None
|
||||
flash_net.eval()
|
||||
|
||||
with amp.autocast(enabled=True):
|
||||
# warm up
|
||||
for i in range(5):
|
||||
y = flash_net(x, context)
|
||||
torch.cuda.synchronize()
|
||||
s1 = time.time()
|
||||
for i in range(10):
|
||||
y = flash_net(x, context)
|
||||
torch.cuda.synchronize()
|
||||
s2 = time.time()
|
||||
|
||||
print(f'Average cost time {(s2-s1)*1000/10} ms')
|
||||
@@ -0,0 +1,2 @@
|
||||
from .clip import *
|
||||
from .midas import *
|
||||
460
modelscope/models/multi_modal/videocomposer/models/clip.py
Normal file
460
modelscope/models/multi_modal/videocomposer/models/clip.py
Normal file
@@ -0,0 +1,460 @@
|
||||
import math
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import modelscope.models.multi_modal.videocomposer.ops as ops
|
||||
|
||||
__all__ = [
|
||||
'CLIP', 'clip_vit_b_32', 'clip_vit_b_16', 'clip_vit_l_14',
|
||||
'clip_vit_l_14_336px', 'clip_vit_h_16'
|
||||
]
|
||||
|
||||
|
||||
def DOWNLOAD_TO_CACHE(oss_key,
|
||||
file_or_dirname=None,
|
||||
cache_dir=osp.join(
|
||||
'/'.join(osp.abspath(__file__).split('/')[:-2]),
|
||||
'model_weights')):
|
||||
r"""Download OSS [file or folder] to the cache folder.
|
||||
Only the 0th process on each node will run the downloading.
|
||||
Barrier all processes until the downloading is completed.
|
||||
"""
|
||||
# source and target paths
|
||||
base_path = osp.join(cache_dir, file_or_dirname or osp.basename(oss_key))
|
||||
|
||||
return base_path
|
||||
|
||||
|
||||
def to_fp16(m):
|
||||
if isinstance(m, (nn.Linear, nn.Conv2d)):
|
||||
m.weight.data = m.weight.data.half()
|
||||
if m.bias is not None:
|
||||
m.bias.data = m.bias.data.half()
|
||||
elif hasattr(m, 'head'):
|
||||
p = getattr(m, 'head')
|
||||
p.data = p.data.half()
|
||||
|
||||
|
||||
class QuickGELU(nn.Module):
|
||||
|
||||
def forward(self, x):
|
||||
return x * torch.sigmoid(1.702 * x)
|
||||
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
r"""Subclass of nn.LayerNorm to handle fp16.
|
||||
"""
|
||||
|
||||
def forward(self, x):
|
||||
return super(LayerNorm, self).forward(x.float()).type_as(x)
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_heads, attn_dropout=0.0, proj_dropout=0.0):
|
||||
assert dim % num_heads == 0
|
||||
super(SelfAttention, self).__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.scale = 1.0 / math.sqrt(self.head_dim)
|
||||
|
||||
# layers
|
||||
self.to_qkv = nn.Linear(dim, dim * 3)
|
||||
self.attn_dropout = nn.Dropout(attn_dropout)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
self.proj_dropout = nn.Dropout(proj_dropout)
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
r"""x: [B, L, C].
|
||||
mask: [*, L, L].
|
||||
"""
|
||||
b, l, _, n = *x.size(), self.num_heads
|
||||
|
||||
# compute query, key, and value
|
||||
q, k, v = self.to_qkv(x.transpose(0, 1)).chunk(3, dim=-1)
|
||||
q = q.reshape(l, b * n, -1).transpose(0, 1)
|
||||
k = k.reshape(l, b * n, -1).transpose(0, 1)
|
||||
v = v.reshape(l, b * n, -1).transpose(0, 1)
|
||||
|
||||
# compute attention
|
||||
attn = self.scale * torch.bmm(q, k.transpose(1, 2))
|
||||
if mask is not None:
|
||||
attn = attn.masked_fill(mask[:, :l, :l] == 0, float('-inf'))
|
||||
attn = F.softmax(attn.float(), dim=-1).type_as(attn)
|
||||
attn = self.attn_dropout(attn)
|
||||
|
||||
# gather context
|
||||
x = torch.bmm(attn, v)
|
||||
x = x.view(b, n, l, -1).transpose(1, 2).reshape(b, l, -1)
|
||||
|
||||
# output
|
||||
x = self.proj(x)
|
||||
x = self.proj_dropout(x)
|
||||
return x
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_heads, attn_dropout=0.0, proj_dropout=0.0):
|
||||
super(AttentionBlock, self).__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
|
||||
# layers
|
||||
self.norm1 = LayerNorm(dim)
|
||||
self.attn = SelfAttention(dim, num_heads, attn_dropout, proj_dropout)
|
||||
self.norm2 = LayerNorm(dim)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(dim, dim * 4), QuickGELU(), nn.Linear(dim * 4, dim),
|
||||
nn.Dropout(proj_dropout))
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
x = x + self.attn(self.norm1(x), mask)
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
|
||||
|
||||
class VisionTransformer(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
image_size=224,
|
||||
patch_size=16,
|
||||
dim=768,
|
||||
out_dim=512,
|
||||
num_heads=12,
|
||||
num_layers=12,
|
||||
attn_dropout=0.0,
|
||||
proj_dropout=0.0,
|
||||
embedding_dropout=0.0):
|
||||
assert image_size % patch_size == 0
|
||||
super(VisionTransformer, self).__init__()
|
||||
self.image_size = image_size
|
||||
self.patch_size = patch_size
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.num_patches = (image_size // patch_size)**2
|
||||
|
||||
# embeddings
|
||||
gain = 1.0 / math.sqrt(dim)
|
||||
self.patch_embedding = nn.Conv2d(
|
||||
3, dim, kernel_size=patch_size, stride=patch_size, bias=False)
|
||||
self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
|
||||
self.pos_embedding = nn.Parameter(
|
||||
gain * torch.randn(1, self.num_patches + 1, dim))
|
||||
self.dropout = nn.Dropout(embedding_dropout)
|
||||
|
||||
# transformer
|
||||
self.pre_norm = LayerNorm(dim)
|
||||
self.transformer = nn.Sequential(*[
|
||||
AttentionBlock(dim, num_heads, attn_dropout, proj_dropout)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
self.post_norm = LayerNorm(dim)
|
||||
|
||||
# head
|
||||
self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
|
||||
|
||||
def forward(self, x):
|
||||
b, dtype = x.size(0), self.head.dtype
|
||||
x = x.type(dtype)
|
||||
|
||||
# patch-embedding
|
||||
x = self.patch_embedding(x).flatten(2).permute(0, 2, 1) # [b, n, c]
|
||||
x = torch.cat([self.cls_embedding.repeat(b, 1, 1).type(dtype), x],
|
||||
dim=1)
|
||||
x = self.dropout(x + self.pos_embedding.type(dtype))
|
||||
x = self.pre_norm(x)
|
||||
|
||||
# transformer
|
||||
x = self.transformer(x)
|
||||
|
||||
# head
|
||||
x = self.post_norm(x)
|
||||
x = torch.mm(x[:, 0, :], self.head)
|
||||
return x
|
||||
|
||||
def fp16(self):
|
||||
return self.apply(to_fp16)
|
||||
|
||||
|
||||
class TextTransformer(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
vocab_size,
|
||||
text_len,
|
||||
dim=512,
|
||||
out_dim=512,
|
||||
num_heads=8,
|
||||
num_layers=12,
|
||||
attn_dropout=0.0,
|
||||
proj_dropout=0.0,
|
||||
embedding_dropout=0.0):
|
||||
super(TextTransformer, self).__init__()
|
||||
self.vocab_size = vocab_size
|
||||
self.text_len = text_len
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
|
||||
# embeddings
|
||||
self.token_embedding = nn.Embedding(vocab_size, dim)
|
||||
self.pos_embedding = nn.Parameter(0.01 * torch.randn(1, text_len, dim))
|
||||
self.dropout = nn.Dropout(embedding_dropout)
|
||||
|
||||
# transformer
|
||||
self.transformer = nn.ModuleList([
|
||||
AttentionBlock(dim, num_heads, attn_dropout, proj_dropout)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
self.norm = LayerNorm(dim)
|
||||
|
||||
# head
|
||||
gain = 1.0 / math.sqrt(dim)
|
||||
self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
|
||||
|
||||
# causal attention mask
|
||||
self.register_buffer('attn_mask',
|
||||
torch.tril(torch.ones(1, text_len, text_len)))
|
||||
|
||||
def forward(self, x):
|
||||
eot, dtype = x.argmax(dim=-1), self.head.dtype
|
||||
|
||||
# embeddings
|
||||
x = self.dropout(
|
||||
self.token_embedding(x).type(dtype)
|
||||
+ self.pos_embedding.type(dtype))
|
||||
|
||||
# transformer
|
||||
for block in self.transformer:
|
||||
x = block(x, self.attn_mask)
|
||||
|
||||
# head
|
||||
x = self.norm(x)
|
||||
x = torch.mm(x[torch.arange(x.size(0)), eot], self.head)
|
||||
return x
|
||||
|
||||
def fp16(self):
|
||||
return self.apply(to_fp16)
|
||||
|
||||
|
||||
class CLIP(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
embed_dim=512,
|
||||
image_size=224,
|
||||
patch_size=16,
|
||||
vision_dim=768,
|
||||
vision_heads=12,
|
||||
vision_layers=12,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=512,
|
||||
text_heads=8,
|
||||
text_layers=12,
|
||||
attn_dropout=0.0,
|
||||
proj_dropout=0.0,
|
||||
embedding_dropout=0.0):
|
||||
super(CLIP, self).__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.image_size = image_size
|
||||
self.patch_size = patch_size
|
||||
self.vision_dim = vision_dim
|
||||
self.vision_heads = vision_heads
|
||||
self.vision_layers = vision_layers
|
||||
self.vocab_size = vocab_size
|
||||
self.text_len = text_len
|
||||
self.text_dim = text_dim
|
||||
self.text_heads = text_heads
|
||||
self.text_layers = text_layers
|
||||
|
||||
# models
|
||||
self.visual = VisionTransformer(
|
||||
image_size=image_size,
|
||||
patch_size=patch_size,
|
||||
dim=vision_dim,
|
||||
out_dim=embed_dim,
|
||||
num_heads=vision_heads,
|
||||
num_layers=vision_layers,
|
||||
attn_dropout=attn_dropout,
|
||||
proj_dropout=proj_dropout,
|
||||
embedding_dropout=embedding_dropout)
|
||||
self.textual = TextTransformer(
|
||||
vocab_size=vocab_size,
|
||||
text_len=text_len,
|
||||
dim=text_dim,
|
||||
out_dim=embed_dim,
|
||||
num_heads=text_heads,
|
||||
num_layers=text_layers,
|
||||
attn_dropout=attn_dropout,
|
||||
proj_dropout=proj_dropout,
|
||||
embedding_dropout=embedding_dropout)
|
||||
self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))
|
||||
|
||||
def forward(self, imgs, txt_tokens):
|
||||
r"""imgs: [B, C, H, W] of torch.float32.
|
||||
txt_tokens: [B, T] of torch.long.
|
||||
"""
|
||||
xi = self.visual(imgs)
|
||||
xt = self.textual(txt_tokens)
|
||||
|
||||
# normalize features
|
||||
xi = F.normalize(xi, p=2, dim=1)
|
||||
xt = F.normalize(xt, p=2, dim=1)
|
||||
|
||||
# gather features from all ranks
|
||||
full_xi = ops.diff_all_gather(xi)
|
||||
full_xt = ops.diff_all_gather(xt)
|
||||
|
||||
# logits
|
||||
scale = self.log_scale.exp()
|
||||
logits_i2t = scale * torch.mm(xi, full_xt.t())
|
||||
logits_t2i = scale * torch.mm(xt, full_xi.t())
|
||||
|
||||
# labels
|
||||
labels = torch.arange(
|
||||
len(xi) * ops.get_rank(),
|
||||
len(xi) * (ops.get_rank() + 1),
|
||||
dtype=torch.long,
|
||||
device=xi.device)
|
||||
return logits_i2t, logits_t2i, labels
|
||||
|
||||
def init_weights(self):
|
||||
# embeddings
|
||||
nn.init.normal_(self.textual.token_embedding.weight, std=0.02)
|
||||
nn.init.normal_(self.visual.patch_embedding.weight, tsd=0.1)
|
||||
|
||||
# attentions
|
||||
for modality in ['visual', 'textual']:
|
||||
dim = self.vision_dim if modality == 'visual' else 'textual'
|
||||
transformer = getattr(self, modality).transformer
|
||||
proj_gain = (1.0 / math.sqrt(dim)) * (
|
||||
1.0 / math.sqrt(2 * transformer.num_layers))
|
||||
attn_gain = 1.0 / math.sqrt(dim)
|
||||
mlp_gain = 1.0 / math.sqrt(2.0 * dim)
|
||||
for block in transformer.layers:
|
||||
nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain)
|
||||
nn.init.normal_(block.attn.proj.weight, std=proj_gain)
|
||||
nn.init.normal_(block.mlp[0].weight, std=mlp_gain)
|
||||
nn.init.normal_(block.mlp[2].weight, std=proj_gain)
|
||||
|
||||
def param_groups(self):
|
||||
groups = [{
|
||||
'params': [
|
||||
p for n, p in self.named_parameters()
|
||||
if 'norm' in n or n.endswith('bias')
|
||||
],
|
||||
'weight_decay':
|
||||
0.0
|
||||
}, {
|
||||
'params': [
|
||||
p for n, p in self.named_parameters()
|
||||
if not ('norm' in n or n.endswith('bias'))
|
||||
]
|
||||
}]
|
||||
return groups
|
||||
|
||||
def fp16(self):
|
||||
return self.apply(to_fp16)
|
||||
|
||||
|
||||
def _clip(name, pretrained=False, **kwargs):
|
||||
model = CLIP(**kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
DOWNLOAD_TO_CACHE(f'models/clip/{name}.pth'),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
|
||||
|
||||
def clip_vit_b_32(pretrained=False, **kwargs):
|
||||
cfg = dict(
|
||||
embed_dim=512,
|
||||
image_size=224,
|
||||
patch_size=32,
|
||||
vision_dim=768,
|
||||
vision_heads=12,
|
||||
vision_layers=12,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=512,
|
||||
text_heads=8,
|
||||
text_layers=12)
|
||||
cfg.update(**kwargs)
|
||||
return _clip('openai-clip-vit-base-32', pretrained, **cfg)
|
||||
|
||||
|
||||
def clip_vit_b_16(pretrained=False, **kwargs):
|
||||
cfg = dict(
|
||||
embed_dim=512,
|
||||
image_size=224,
|
||||
patch_size=32,
|
||||
vision_dim=768,
|
||||
vision_heads=12,
|
||||
vision_layers=12,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=512,
|
||||
text_heads=8,
|
||||
text_layers=12)
|
||||
cfg.update(**kwargs)
|
||||
return _clip('openai-clip-vit-base-16', pretrained, **cfg)
|
||||
|
||||
|
||||
def clip_vit_l_14(pretrained=False, **kwargs):
|
||||
cfg = dict(
|
||||
embed_dim=768,
|
||||
image_size=224,
|
||||
patch_size=14,
|
||||
vision_dim=1024,
|
||||
vision_heads=16,
|
||||
vision_layers=24,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=768,
|
||||
text_heads=12,
|
||||
text_layers=12)
|
||||
cfg.update(**kwargs)
|
||||
return _clip('openai-clip-vit-large-14', pretrained, **cfg)
|
||||
|
||||
|
||||
def clip_vit_l_14_336px(pretrained=False, **kwargs):
|
||||
cfg = dict(
|
||||
embed_dim=768,
|
||||
image_size=336,
|
||||
patch_size=14,
|
||||
vision_dim=1024,
|
||||
vision_heads=16,
|
||||
vision_layers=24,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=768,
|
||||
text_heads=12,
|
||||
text_layers=12)
|
||||
cfg.update(**kwargs)
|
||||
return _clip('openai-clip-vit-large-14-336px', pretrained, **cfg)
|
||||
|
||||
|
||||
def clip_vit_h_16(pretrained=False, **kwargs):
|
||||
assert not pretrained, 'pretrained model for openai-clip-vit-huge-16 is not available!'
|
||||
cfg = dict(
|
||||
embed_dim=1024,
|
||||
image_size=256,
|
||||
patch_size=16,
|
||||
vision_dim=1280,
|
||||
vision_heads=16,
|
||||
vision_layers=32,
|
||||
vocab_size=49408,
|
||||
text_len=77,
|
||||
text_dim=1024,
|
||||
text_heads=16,
|
||||
text_layers=24)
|
||||
cfg.update(**kwargs)
|
||||
return _clip('openai-clip-vit-huge-16', pretrained, **cfg)
|
||||
320
modelscope/models/multi_modal/videocomposer/models/midas.py
Normal file
320
modelscope/models/multi_modal/videocomposer/models/midas.py
Normal file
@@ -0,0 +1,320 @@
|
||||
r"""A much cleaner re-implementation of ``https://github.com/isl-org/MiDaS''.
|
||||
Image augmentation: T.Compose([
|
||||
Resize(
|
||||
keep_aspect_ratio=True,
|
||||
ensure_multiple_of=32,
|
||||
interpolation=cv2.INTER_CUBIC),
|
||||
T.ToTensor(),
|
||||
T.Normalize(
|
||||
mean=[0.5, 0.5, 0.5],
|
||||
std=[0.5, 0.5, 0.5])]).
|
||||
Fast inference:
|
||||
model = model.to(memory_format=torch.channels_last).half()
|
||||
input = input.to(memory_format=torch.channels_last).half()
|
||||
output = model(input)
|
||||
"""
|
||||
import math
|
||||
import os
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
__all__ = ['MiDaS', 'midas_v3']
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_heads):
|
||||
assert dim % num_heads == 0
|
||||
super(SelfAttention, self).__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.scale = 1.0 / math.sqrt(self.head_dim)
|
||||
|
||||
# layers
|
||||
self.to_qkv = nn.Linear(dim, dim * 3)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
def forward(self, x):
|
||||
b, l, c, n, d = *x.size(), self.num_heads, self.head_dim
|
||||
|
||||
# compute query, key, value
|
||||
q, k, v = self.to_qkv(x).view(b, l, n * 3, d).chunk(3, dim=2)
|
||||
|
||||
# compute attention
|
||||
attn = self.scale * torch.einsum('binc,bjnc->bnij', q, k)
|
||||
attn = F.softmax(attn.float(), dim=-1).type_as(attn)
|
||||
|
||||
# gather context
|
||||
x = torch.einsum('bnij,bjnc->binc', attn, v)
|
||||
x = x.reshape(b, l, c)
|
||||
|
||||
# output
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_heads):
|
||||
super(AttentionBlock, self).__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
|
||||
# layers
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.attn = SelfAttention(dim, num_heads)
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim))
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.attn(self.norm1(x))
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
|
||||
|
||||
class VisionTransformer(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
image_size=384,
|
||||
patch_size=16,
|
||||
dim=1024,
|
||||
out_dim=1000,
|
||||
num_heads=16,
|
||||
num_layers=24):
|
||||
assert image_size % patch_size == 0
|
||||
super(VisionTransformer, self).__init__()
|
||||
self.image_size = image_size
|
||||
self.patch_size = patch_size
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.num_patches = (image_size // patch_size)**2
|
||||
|
||||
# embeddings
|
||||
self.patch_embedding = nn.Conv2d(
|
||||
3, dim, kernel_size=patch_size, stride=patch_size)
|
||||
self.cls_embedding = nn.Parameter(torch.zeros(1, 1, dim))
|
||||
self.pos_embedding = nn.Parameter(
|
||||
torch.empty(1, self.num_patches + 1, dim).normal_(std=0.02))
|
||||
|
||||
# blocks
|
||||
self.blocks = nn.Sequential(
|
||||
*[AttentionBlock(dim, num_heads) for _ in range(num_layers)])
|
||||
self.norm = nn.LayerNorm(dim)
|
||||
|
||||
# head
|
||||
self.head = nn.Linear(dim, out_dim)
|
||||
|
||||
def forward(self, x):
|
||||
b = x.size(0)
|
||||
|
||||
# embeddings
|
||||
x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
|
||||
x = torch.cat([self.cls_embedding.repeat(b, 1, 1), x], dim=1)
|
||||
x = x + self.pos_embedding
|
||||
|
||||
# blocks
|
||||
x = self.blocks(x)
|
||||
x = self.norm(x)
|
||||
|
||||
# head
|
||||
x = self.head(x)
|
||||
return x
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
|
||||
def __init__(self, dim):
|
||||
super(ResidualBlock, self).__init__()
|
||||
self.dim = dim
|
||||
|
||||
# layers
|
||||
self.residual = nn.Sequential(
|
||||
nn.ReLU(inplace=False), # NOTE: avoid modifying the input
|
||||
nn.Conv2d(dim, dim, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(dim, dim, 3, padding=1))
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.residual(x)
|
||||
|
||||
|
||||
class FusionBlock(nn.Module):
|
||||
|
||||
def __init__(self, dim):
|
||||
super(FusionBlock, self).__init__()
|
||||
self.dim = dim
|
||||
|
||||
# layers
|
||||
self.layer1 = ResidualBlock(dim)
|
||||
self.layer2 = ResidualBlock(dim)
|
||||
self.conv_out = nn.Conv2d(dim, dim, 1)
|
||||
|
||||
def forward(self, *xs):
|
||||
assert len(xs) in (1, 2), 'invalid number of inputs'
|
||||
if len(xs) == 1:
|
||||
x = self.layer2(xs[0])
|
||||
else:
|
||||
x = self.layer2(xs[0] + self.layer1(xs[1]))
|
||||
x = F.interpolate(
|
||||
x, scale_factor=2, mode='bilinear', align_corners=True)
|
||||
x = self.conv_out(x)
|
||||
return x
|
||||
|
||||
|
||||
class MiDaS(nn.Module):
|
||||
r"""MiDaS v3.0 DPT-Large from ``https://github.com/isl-org/MiDaS''.
|
||||
Monocular depth estimation using dense prediction transformers.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
image_size=384,
|
||||
patch_size=16,
|
||||
dim=1024,
|
||||
neck_dims=[256, 512, 1024, 1024],
|
||||
fusion_dim=256,
|
||||
num_heads=16,
|
||||
num_layers=24):
|
||||
assert image_size % patch_size == 0
|
||||
assert num_layers % 4 == 0
|
||||
super(MiDaS, self).__init__()
|
||||
self.image_size = image_size
|
||||
self.patch_size = patch_size
|
||||
self.dim = dim
|
||||
self.neck_dims = neck_dims
|
||||
self.fusion_dim = fusion_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.num_patches = (image_size // patch_size)**2
|
||||
|
||||
# embeddings
|
||||
self.patch_embedding = nn.Conv2d(
|
||||
3, dim, kernel_size=patch_size, stride=patch_size)
|
||||
self.cls_embedding = nn.Parameter(torch.zeros(1, 1, dim))
|
||||
self.pos_embedding = nn.Parameter(
|
||||
torch.empty(1, self.num_patches + 1, dim).normal_(std=0.02))
|
||||
|
||||
# blocks
|
||||
stride = num_layers // 4
|
||||
self.blocks = nn.Sequential(
|
||||
*[AttentionBlock(dim, num_heads) for _ in range(num_layers)])
|
||||
self.slices = [slice(i * stride, (i + 1) * stride) for i in range(4)]
|
||||
|
||||
# stage1 (4x)
|
||||
self.fc1 = nn.Sequential(nn.Linear(dim * 2, dim), nn.GELU())
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(dim, neck_dims[0], 1),
|
||||
nn.ConvTranspose2d(neck_dims[0], neck_dims[0], 4, stride=4),
|
||||
nn.Conv2d(neck_dims[0], fusion_dim, 3, padding=1, bias=False))
|
||||
self.fusion1 = FusionBlock(fusion_dim)
|
||||
|
||||
# stage2 (8x)
|
||||
self.fc2 = nn.Sequential(nn.Linear(dim * 2, dim), nn.GELU())
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv2d(dim, neck_dims[1], 1),
|
||||
nn.ConvTranspose2d(neck_dims[1], neck_dims[1], 2, stride=2),
|
||||
nn.Conv2d(neck_dims[1], fusion_dim, 3, padding=1, bias=False))
|
||||
self.fusion2 = FusionBlock(fusion_dim)
|
||||
|
||||
# stage3 (16x)
|
||||
self.fc3 = nn.Sequential(nn.Linear(dim * 2, dim), nn.GELU())
|
||||
self.conv3 = nn.Sequential(
|
||||
nn.Conv2d(dim, neck_dims[2], 1),
|
||||
nn.Conv2d(neck_dims[2], fusion_dim, 3, padding=1, bias=False))
|
||||
self.fusion3 = FusionBlock(fusion_dim)
|
||||
|
||||
# stage4 (32x)
|
||||
self.fc4 = nn.Sequential(nn.Linear(dim * 2, dim), nn.GELU())
|
||||
self.conv4 = nn.Sequential(
|
||||
nn.Conv2d(dim, neck_dims[3], 1),
|
||||
nn.Conv2d(neck_dims[3], neck_dims[3], 3, stride=2, padding=1),
|
||||
nn.Conv2d(neck_dims[3], fusion_dim, 3, padding=1, bias=False))
|
||||
self.fusion4 = FusionBlock(fusion_dim)
|
||||
|
||||
# head
|
||||
self.head = nn.Sequential(
|
||||
nn.Conv2d(fusion_dim, fusion_dim // 2, 3, padding=1),
|
||||
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
|
||||
nn.Conv2d(fusion_dim // 2, 32, 3, padding=1),
|
||||
nn.ReLU(inplace=True), nn.ConvTranspose2d(32, 1, 1),
|
||||
nn.ReLU(inplace=True))
|
||||
|
||||
def forward(self, x):
|
||||
b, _, h, w, p = *x.size(), self.patch_size
|
||||
assert h % p == 0 and w % p == 0, f'Image size ({w}, {h}) is not divisible by patch size ({p}, {p})'
|
||||
hp, wp, grid = h // p, w // p, self.image_size // p
|
||||
|
||||
# embeddings
|
||||
pos_embedding = torch.cat([
|
||||
self.pos_embedding[:, :1],
|
||||
F.interpolate(
|
||||
self.pos_embedding[:, 1:].reshape(1, grid, grid, -1).permute(
|
||||
0, 3, 1, 2),
|
||||
size=(hp, wp),
|
||||
mode='bilinear',
|
||||
align_corners=False).permute(0, 2, 3, 1).reshape(
|
||||
1, hp * wp, -1)
|
||||
],
|
||||
dim=1) # noqa
|
||||
x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
|
||||
x = torch.cat([self.cls_embedding.repeat(b, 1, 1), x], dim=1)
|
||||
x = x + pos_embedding
|
||||
|
||||
# stage1
|
||||
x = self.blocks[self.slices[0]](x)
|
||||
x1 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
|
||||
x1 = self.fc1(x1).permute(0, 2, 1).unflatten(2, (hp, wp))
|
||||
x1 = self.conv1(x1)
|
||||
|
||||
# stage2
|
||||
x = self.blocks[self.slices[1]](x)
|
||||
x2 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
|
||||
x2 = self.fc2(x2).permute(0, 2, 1).unflatten(2, (hp, wp))
|
||||
x2 = self.conv2(x2)
|
||||
|
||||
# stage3
|
||||
x = self.blocks[self.slices[2]](x)
|
||||
x3 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
|
||||
x3 = self.fc3(x3).permute(0, 2, 1).unflatten(2, (hp, wp))
|
||||
x3 = self.conv3(x3)
|
||||
|
||||
# stage4
|
||||
x = self.blocks[self.slices[3]](x)
|
||||
x4 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
|
||||
x4 = self.fc4(x4).permute(0, 2, 1).unflatten(2, (hp, wp))
|
||||
x4 = self.conv4(x4)
|
||||
|
||||
# fusion
|
||||
x4 = self.fusion4(x4)
|
||||
x3 = self.fusion3(x4, x3)
|
||||
x2 = self.fusion2(x3, x2)
|
||||
x1 = self.fusion1(x2, x1)
|
||||
|
||||
# head
|
||||
x = self.head(x1)
|
||||
return x
|
||||
|
||||
|
||||
def midas_v3(model_dir, pretrained=False, **kwargs):
|
||||
cfg = dict(
|
||||
image_size=384,
|
||||
patch_size=16,
|
||||
dim=1024,
|
||||
neck_dims=[256, 512, 1024, 1024],
|
||||
fusion_dim=256,
|
||||
num_heads=16,
|
||||
num_layers=24)
|
||||
cfg.update(**kwargs)
|
||||
model = MiDaS(**cfg)
|
||||
if pretrained:
|
||||
model.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(model_dir, 'midas_v3_dpt_large.pth'),
|
||||
map_location='cpu'))
|
||||
return model
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from .degration import *
|
||||
from .distributed import *
|
||||
from .losses import *
|
||||
from .random_mask import *
|
||||
from .utils import *
|
||||
998
modelscope/models/multi_modal/videocomposer/ops/degration.py
Normal file
998
modelscope/models/multi_modal/videocomposer/ops/degration.py
Normal file
@@ -0,0 +1,998 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
import scipy
|
||||
import scipy.stats as stats
|
||||
import torch
|
||||
from scipy import ndimage
|
||||
from scipy.interpolate import interp2d
|
||||
from scipy.linalg import orth
|
||||
from torchvision.utils import make_grid
|
||||
|
||||
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
|
||||
|
||||
__all__ = ['degradation_bsrgan_light', 'degradation_bsrgan']
|
||||
|
||||
|
||||
# get uint8 image of size HxWxn_channles (RGB)
|
||||
def imread_uint(path, n_channels=3):
|
||||
# input: path
|
||||
# output: HxWx3(RGB or GGG), or HxWx1 (G)
|
||||
if n_channels == 1:
|
||||
img = cv2.imread(path, 0)
|
||||
img = np.expand_dims(img, axis=2)
|
||||
elif n_channels == 3:
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
if img.ndim == 2:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
||||
else:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
return img
|
||||
|
||||
|
||||
def uint2single(img):
|
||||
return np.float32(img / 255.)
|
||||
|
||||
|
||||
def single2uint(img):
|
||||
return np.uint8((img.clip(0, 1) * 255.).round())
|
||||
|
||||
|
||||
def uint162single(img):
|
||||
return np.float32(img / 65535.)
|
||||
|
||||
|
||||
def single2uint16(img):
|
||||
return np.uint16((img.clip(0, 1) * 65535.).round())
|
||||
|
||||
|
||||
def rgb2ycbcr(img, only_y=True):
|
||||
'''same as matlab rgb2ycbcr
|
||||
only_y: only return Y channel
|
||||
Input:
|
||||
uint8, [0, 255]
|
||||
float, [0, 1]
|
||||
'''
|
||||
in_img_type = img.dtype
|
||||
img.astype(np.float32)
|
||||
if in_img_type != np.uint8:
|
||||
img *= 255.
|
||||
# convert
|
||||
if only_y:
|
||||
rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
|
||||
else:
|
||||
rlt = np.matmul(img,
|
||||
[[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
|
||||
[24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
|
||||
if in_img_type == np.uint8:
|
||||
rlt = rlt.round()
|
||||
else:
|
||||
rlt /= 255.
|
||||
return rlt.astype(in_img_type)
|
||||
|
||||
|
||||
def ycbcr2rgb(img):
|
||||
'''same as matlab ycbcr2rgb
|
||||
Input:
|
||||
uint8, [0, 255]
|
||||
float, [0, 1]
|
||||
'''
|
||||
in_img_type = img.dtype
|
||||
img.astype(np.float32)
|
||||
if in_img_type != np.uint8:
|
||||
img *= 255.
|
||||
# convert
|
||||
rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621],
|
||||
[0, -0.00153632, 0.00791071],
|
||||
[0.00625893, -0.00318811, 0]]) * 255.0 + [
|
||||
-222.921, 135.576, -276.836
|
||||
] # noqa
|
||||
if in_img_type == np.uint8:
|
||||
rlt = rlt.round()
|
||||
else:
|
||||
rlt /= 255.
|
||||
return rlt.astype(in_img_type)
|
||||
|
||||
|
||||
def bgr2ycbcr(img, only_y=True):
|
||||
'''bgr version of rgb2ycbcr
|
||||
only_y: only return Y channel
|
||||
Input:
|
||||
uint8, [0, 255]
|
||||
float, [0, 1]
|
||||
'''
|
||||
in_img_type = img.dtype
|
||||
img.astype(np.float32)
|
||||
if in_img_type != np.uint8:
|
||||
img *= 255.
|
||||
# convert
|
||||
if only_y:
|
||||
rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
|
||||
else:
|
||||
rlt = np.matmul(img,
|
||||
[[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
|
||||
[65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
|
||||
if in_img_type == np.uint8:
|
||||
rlt = rlt.round()
|
||||
else:
|
||||
rlt /= 255.
|
||||
return rlt.astype(in_img_type)
|
||||
|
||||
|
||||
def channel_convert(in_c, tar_type, img_list):
|
||||
# conversion among BGR, gray and y
|
||||
if in_c == 3 and tar_type == 'gray':
|
||||
gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
|
||||
return [np.expand_dims(img, axis=2) for img in gray_list]
|
||||
elif in_c == 3 and tar_type == 'y':
|
||||
y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
|
||||
return [np.expand_dims(img, axis=2) for img in y_list]
|
||||
elif in_c == 1 and tar_type == 'RGB':
|
||||
return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
|
||||
else:
|
||||
return img_list
|
||||
|
||||
|
||||
# PSNR
|
||||
def calculate_psnr(img1, img2, border=0):
|
||||
# img1 and img2 have range [0, 255]
|
||||
if not img1.shape == img2.shape:
|
||||
raise ValueError('Input images must have the same dimensions.')
|
||||
h, w = img1.shape[:2]
|
||||
img1 = img1[border:h - border, border:w - border]
|
||||
img2 = img2[border:h - border, border:w - border]
|
||||
|
||||
img1 = img1.astype(np.float64)
|
||||
img2 = img2.astype(np.float64)
|
||||
mse = np.mean((img1 - img2)**2)
|
||||
if mse == 0:
|
||||
return float('inf')
|
||||
return 20 * math.log10(255.0 / math.sqrt(mse))
|
||||
|
||||
|
||||
# SSIM
|
||||
def calculate_ssim(img1, img2, border=0):
|
||||
'''calculate SSIM
|
||||
the same outputs as MATLAB's
|
||||
img1, img2: [0, 255]
|
||||
'''
|
||||
if not img1.shape == img2.shape:
|
||||
raise ValueError('Input images must have the same dimensions.')
|
||||
h, w = img1.shape[:2]
|
||||
img1 = img1[border:h - border, border:w - border]
|
||||
img2 = img2[border:h - border, border:w - border]
|
||||
|
||||
if img1.ndim == 2:
|
||||
return ssim(img1, img2)
|
||||
elif img1.ndim == 3:
|
||||
if img1.shape[2] == 3:
|
||||
ssims = []
|
||||
for i in range(3):
|
||||
ssims.append(ssim(img1[:, :, i], img2[:, :, i]))
|
||||
return np.array(ssims).mean()
|
||||
elif img1.shape[2] == 1:
|
||||
return ssim(np.squeeze(img1), np.squeeze(img2))
|
||||
else:
|
||||
raise ValueError('Wrong input image dimensions.')
|
||||
|
||||
|
||||
def ssim(img1, img2):
|
||||
C1 = (0.01 * 255)**2
|
||||
C2 = (0.03 * 255)**2
|
||||
|
||||
img1 = img1.astype(np.float64)
|
||||
img2 = img2.astype(np.float64)
|
||||
kernel = cv2.getGaussianKernel(11, 1.5)
|
||||
window = np.outer(kernel, kernel.transpose())
|
||||
|
||||
mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
|
||||
mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
|
||||
mu1_sq = mu1**2
|
||||
mu2_sq = mu2**2
|
||||
mu1_mu2 = mu1 * mu2
|
||||
sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
|
||||
sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
|
||||
sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
|
||||
|
||||
ssim_map = ((2 * mu1_mu2 + C1) # noqa
|
||||
* (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) # noqa
|
||||
* # noqa
|
||||
(sigma1_sq + sigma2_sq + C2)) # noqa
|
||||
return ssim_map.mean()
|
||||
|
||||
|
||||
# matlab 'imresize' function, now only support 'bicubic'
|
||||
def cubic(x):
|
||||
absx = torch.abs(x)
|
||||
absx2 = absx**2
|
||||
absx3 = absx**3
|
||||
return (1.5 * absx3 - 2.5 * absx2 + 1) * ((absx <= 1).type_as(absx)) + \
|
||||
(-0.5 * absx3 + 2.5 * absx2 - 4*absx + 2) * (((absx > 1) * (absx <= 2)).type_as(absx)) # noqa
|
||||
|
||||
|
||||
def calculate_weights_indices(in_length, out_length, scale, kernel,
|
||||
kernel_width, antialiasing):
|
||||
if (scale < 1) and (antialiasing):
|
||||
# Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
|
||||
kernel_width = kernel_width / scale
|
||||
|
||||
# Output-space coordinates
|
||||
x = torch.linspace(1, out_length, out_length)
|
||||
|
||||
# Input-space coordinates. Calculate the inverse mapping such that 0.5
|
||||
# in output space maps to 0.5 in input space, and 0.5+scale in output
|
||||
# space maps to 1.5 in input space.
|
||||
u = x / scale + 0.5 * (1 - 1 / scale)
|
||||
|
||||
# What is the left-most pixel that can be involved in the computation?
|
||||
left = torch.floor(u - kernel_width / 2)
|
||||
|
||||
# What is the maximum number of pixels that can be involved in the
|
||||
# computation? Note: it's OK to use an extra pixel here; if the
|
||||
# corresponding weights are all zero, it will be eliminated at the end
|
||||
# of this function.
|
||||
P = math.ceil(kernel_width) + 2
|
||||
|
||||
# The indices of the input pixels involved in computing the k-th output
|
||||
# pixel are in row k of the indices matrix.
|
||||
indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(
|
||||
0, P - 1, P).view(1, P).expand(out_length, P)
|
||||
|
||||
# The weights used to compute the k-th output pixel are in row k of the
|
||||
# weights matrix.
|
||||
distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
|
||||
# apply cubic kernel
|
||||
if (scale < 1) and (antialiasing):
|
||||
weights = scale * cubic(distance_to_center * scale)
|
||||
else:
|
||||
weights = cubic(distance_to_center)
|
||||
# Normalize the weights matrix so that each row sums to 1.
|
||||
weights_sum = torch.sum(weights, 1).view(out_length, 1)
|
||||
weights = weights / weights_sum.expand(out_length, P)
|
||||
|
||||
# If a column in weights is all zero, get rid of it. only consider the first and last column.
|
||||
weights_zero_tmp = torch.sum((weights == 0), 0)
|
||||
if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
|
||||
indices = indices.narrow(1, 1, P - 2)
|
||||
weights = weights.narrow(1, 1, P - 2)
|
||||
if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
|
||||
indices = indices.narrow(1, 0, P - 2)
|
||||
weights = weights.narrow(1, 0, P - 2)
|
||||
weights = weights.contiguous()
|
||||
indices = indices.contiguous()
|
||||
sym_len_s = -indices.min() + 1
|
||||
sym_len_e = indices.max() - in_length
|
||||
indices = indices + sym_len_s - 1
|
||||
return weights, indices, int(sym_len_s), int(sym_len_e)
|
||||
|
||||
|
||||
# imresize for tensor image [0, 1]
|
||||
def imresize(img, scale, antialiasing=True):
|
||||
# Now the scale should be the same for H and W
|
||||
# input: img: pytorch tensor, CHW or HW [0,1]
|
||||
# output: CHW or HW [0,1] w/o round
|
||||
need_squeeze = True if img.dim() == 2 else False
|
||||
if need_squeeze:
|
||||
img.unsqueeze_(0)
|
||||
in_C, in_H, in_W = img.size()
|
||||
out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W
|
||||
* scale)
|
||||
kernel_width = 4
|
||||
kernel = 'cubic'
|
||||
|
||||
# Return the desired dimension order for performing the resize. The
|
||||
# strategy is to perform the resize first along the dimension with the
|
||||
# smallest scale factor.
|
||||
# Now we do not support this.
|
||||
|
||||
# get weights and indices
|
||||
weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
|
||||
in_H, out_H, scale, kernel, kernel_width, antialiasing)
|
||||
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
|
||||
in_W, out_W, scale, kernel, kernel_width, antialiasing)
|
||||
# process H dimension
|
||||
# symmetric copying
|
||||
img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
|
||||
img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
|
||||
|
||||
sym_patch = img[:, :sym_len_Hs, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = img[:, -sym_len_He:, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
|
||||
|
||||
out_1 = torch.FloatTensor(in_C, out_H, in_W)
|
||||
kernel_width = weights_H.size(1)
|
||||
for i in range(out_H):
|
||||
idx = int(indices_H[i][0])
|
||||
for j in range(out_C):
|
||||
out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(
|
||||
0, 1).mv(weights_H[i])
|
||||
|
||||
# process W dimension
|
||||
# symmetric copying
|
||||
out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
|
||||
out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
|
||||
|
||||
sym_patch = out_1[:, :, :sym_len_Ws]
|
||||
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
||||
out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = out_1[:, :, -sym_len_We:]
|
||||
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
||||
out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
|
||||
|
||||
out_2 = torch.FloatTensor(in_C, out_H, out_W)
|
||||
kernel_width = weights_W.size(1)
|
||||
for i in range(out_W):
|
||||
idx = int(indices_W[i][0])
|
||||
for j in range(out_C):
|
||||
out_2[j, :, i] = out_1_aug[j, :,
|
||||
idx:idx + kernel_width].mv(weights_W[i])
|
||||
if need_squeeze:
|
||||
out_2.squeeze_()
|
||||
return out_2
|
||||
|
||||
|
||||
# imresize for numpy image [0, 1]
|
||||
def imresize_np(img, scale, antialiasing=True):
|
||||
# Now the scale should be the same for H and W
|
||||
# input: img: Numpy, HWC or HW [0,1]
|
||||
# output: HWC or HW [0,1] w/o round
|
||||
img = torch.from_numpy(img)
|
||||
need_squeeze = True if img.dim() == 2 else False
|
||||
if need_squeeze:
|
||||
img.unsqueeze_(2)
|
||||
|
||||
in_H, in_W, in_C = img.size()
|
||||
out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W
|
||||
* scale)
|
||||
kernel_width = 4
|
||||
kernel = 'cubic'
|
||||
|
||||
# Return the desired dimension order for performing the resize. The
|
||||
# strategy is to perform the resize first along the dimension with the
|
||||
# smallest scale factor.
|
||||
# Now we do not support this.
|
||||
|
||||
# get weights and indices
|
||||
weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
|
||||
in_H, out_H, scale, kernel, kernel_width, antialiasing)
|
||||
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
|
||||
in_W, out_W, scale, kernel, kernel_width, antialiasing)
|
||||
# process H dimension
|
||||
# symmetric copying
|
||||
img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
|
||||
img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
|
||||
|
||||
sym_patch = img[:sym_len_Hs, :, :]
|
||||
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
||||
img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = img[-sym_len_He:, :, :]
|
||||
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
||||
img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
|
||||
|
||||
out_1 = torch.FloatTensor(out_H, in_W, in_C)
|
||||
kernel_width = weights_H.size(1)
|
||||
for i in range(out_H):
|
||||
idx = int(indices_H[i][0])
|
||||
for j in range(out_C):
|
||||
out_1[i, :, j] = img_aug[idx:idx + kernel_width, :,
|
||||
j].transpose(0, 1).mv(weights_H[i])
|
||||
|
||||
# process W dimension
|
||||
# symmetric copying
|
||||
out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
|
||||
out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
|
||||
|
||||
sym_patch = out_1[:, :sym_len_Ws, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = out_1[:, -sym_len_We:, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
|
||||
|
||||
out_2 = torch.FloatTensor(out_H, out_W, in_C)
|
||||
kernel_width = weights_W.size(1)
|
||||
for i in range(out_W):
|
||||
idx = int(indices_W[i][0])
|
||||
for j in range(out_C):
|
||||
out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width,
|
||||
j].mv(weights_W[i])
|
||||
if need_squeeze:
|
||||
out_2.squeeze_()
|
||||
|
||||
return out_2.numpy()
|
||||
|
||||
|
||||
def modcrop_np(img, sf):
|
||||
'''
|
||||
Args:
|
||||
img: numpy image, WxH or WxHxC
|
||||
sf: scale factor
|
||||
Return:
|
||||
cropped image
|
||||
'''
|
||||
w, h = img.shape[:2]
|
||||
im = np.copy(img)
|
||||
return im[:w - w % sf, :h - h % sf, ...]
|
||||
|
||||
|
||||
def analytic_kernel(k):
|
||||
"""Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
|
||||
k_size = k.shape[0]
|
||||
# Calculate the big kernels size
|
||||
big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
|
||||
# Loop over the small kernel to fill the big one
|
||||
for r in range(k_size):
|
||||
for c in range(k_size):
|
||||
big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
|
||||
# Crop the edges of the big kernel to ignore very small values and increase run time of SR
|
||||
crop = k_size // 2
|
||||
cropped_big_k = big_k[crop:-crop, crop:-crop]
|
||||
# Normalize to 1
|
||||
return cropped_big_k / cropped_big_k.sum()
|
||||
|
||||
|
||||
def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
|
||||
""" generate an anisotropic Gaussian kernel
|
||||
Args:
|
||||
ksize : e.g., 15, kernel size
|
||||
theta : [0, pi], rotation angle range
|
||||
l1 : [0.1,50], scaling of eigenvalues
|
||||
l2 : [0.1,l1], scaling of eigenvalues
|
||||
If l1 = l2, will get an isotropic Gaussian kernel.
|
||||
Returns:
|
||||
k : kernel
|
||||
"""
|
||||
|
||||
v = np.dot(
|
||||
np.array([[np.cos(theta), -np.sin(theta)],
|
||||
[np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
|
||||
V = np.array([[v[0], v[1]], [v[1], -v[0]]])
|
||||
D = np.array([[l1, 0], [0, l2]])
|
||||
Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
|
||||
k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
|
||||
return k
|
||||
|
||||
|
||||
def gm_blur_kernel(mean, cov, size=15):
|
||||
center = size / 2.0 + 0.5
|
||||
k = np.zeros([size, size])
|
||||
for y in range(size):
|
||||
for x in range(size):
|
||||
cy = y - center + 1
|
||||
cx = x - center + 1
|
||||
k[y, x] = stats.multivariate_normal.pdf([cx, cy],
|
||||
mean=mean,
|
||||
cov=cov)
|
||||
|
||||
k = k / np.sum(k)
|
||||
return k
|
||||
|
||||
|
||||
def shift_pixel(x, sf, upper_left=True):
|
||||
"""shift pixel for super-resolution with different scale factors
|
||||
Args:
|
||||
x: WxHxC or WxH
|
||||
sf: scale factor
|
||||
upper_left: shift direction
|
||||
"""
|
||||
h, w = x.shape[:2]
|
||||
shift = (sf - 1) * 0.5
|
||||
xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
|
||||
if upper_left:
|
||||
x1 = xv + shift
|
||||
y1 = yv + shift
|
||||
else:
|
||||
x1 = xv - shift
|
||||
y1 = yv - shift
|
||||
|
||||
x1 = np.clip(x1, 0, w - 1)
|
||||
y1 = np.clip(y1, 0, h - 1)
|
||||
|
||||
if x.ndim == 2:
|
||||
x = interp2d(xv, yv, x)(x1, y1)
|
||||
if x.ndim == 3:
|
||||
for i in range(x.shape[-1]):
|
||||
x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def blur(x, k):
|
||||
'''
|
||||
x: image, NxcxHxW
|
||||
k: kernel, Nx1xhxw
|
||||
'''
|
||||
n, c = x.shape[:2]
|
||||
p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
|
||||
x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
|
||||
k = k.repeat(1, c, 1, 1)
|
||||
k = k.view(-1, 1, k.shape[2], k.shape[3])
|
||||
x = x.view(1, -1, x.shape[2], x.shape[3])
|
||||
x = torch.nn.functional.conv2d(
|
||||
x, k, bias=None, stride=1, padding=0, groups=n * c)
|
||||
x = x.view(n, c, x.shape[2], x.shape[3])
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def gen_kernel(
|
||||
k_size=np.array([15, 15]),
|
||||
scale_factor=np.array([4, 4]),
|
||||
min_var=0.6,
|
||||
max_var=10.,
|
||||
noise_level=0):
|
||||
""""
|
||||
# modified version of https://github.com/assafshocher/BlindSR_dataset_generator
|
||||
# Kai Zhang
|
||||
# min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
|
||||
# max_var = 2.5 * sf
|
||||
"""
|
||||
# Set random eigen-vals (lambdas) and angle (theta) for COV matrix
|
||||
lambda_1 = min_var + np.random.rand() * (max_var - min_var)
|
||||
lambda_2 = min_var + np.random.rand() * (max_var - min_var)
|
||||
theta = np.random.rand() * np.pi # random theta
|
||||
noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
|
||||
|
||||
# Set COV matrix using Lambdas and Theta
|
||||
LAMBDA = np.diag([lambda_1, lambda_2])
|
||||
Q = np.array([[np.cos(theta), -np.sin(theta)],
|
||||
[np.sin(theta), np.cos(theta)]])
|
||||
SIGMA = Q @ LAMBDA @ Q.T
|
||||
INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
|
||||
|
||||
# Set expectation position (shifting kernel for aligned image)
|
||||
MU = k_size // 2 - 0.5 * (scale_factor - 1)
|
||||
MU = MU[None, None, :, None]
|
||||
|
||||
# Create meshgrid for Gaussian
|
||||
[X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
|
||||
Z = np.stack([X, Y], 2)[:, :, :, None]
|
||||
|
||||
# Calcualte Gaussian for every pixel of the kernel
|
||||
ZZ = Z - MU
|
||||
ZZ_t = ZZ.transpose(0, 1, 3, 2)
|
||||
raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
|
||||
|
||||
# shift the kernel so it will be centered
|
||||
# raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
|
||||
|
||||
# Normalize the kernel and return
|
||||
# kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
|
||||
kernel = raw_kernel / np.sum(raw_kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def fspecial_gaussian(hsize, sigma):
|
||||
hsize = [hsize, hsize]
|
||||
siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
|
||||
std = sigma
|
||||
[x, y] = np.meshgrid(
|
||||
np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
|
||||
arg = -(x * x + y * y) / (2 * std * std)
|
||||
h = np.exp(arg)
|
||||
h[h < scipy.finfo(float).eps * h.max()] = 0
|
||||
sumh = h.sum()
|
||||
if sumh != 0:
|
||||
h = h / sumh
|
||||
return h
|
||||
|
||||
|
||||
def fspecial_laplacian(alpha):
|
||||
alpha = max([0, min([alpha, 1])])
|
||||
h1 = alpha / (alpha + 1)
|
||||
h2 = (1 - alpha) / (alpha + 1)
|
||||
h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
|
||||
h = np.array(h)
|
||||
return h
|
||||
|
||||
|
||||
def fspecial(filter_type, *args, **kwargs):
|
||||
if filter_type == 'gaussian':
|
||||
return fspecial_gaussian(*args, **kwargs)
|
||||
if filter_type == 'laplacian':
|
||||
return fspecial_laplacian(*args, **kwargs)
|
||||
|
||||
|
||||
def bicubic_degradation(x, sf=3):
|
||||
'''
|
||||
Args:
|
||||
x: HxWxC image, [0, 1]
|
||||
sf: down-scale factor
|
||||
Return:
|
||||
bicubicly downsampled LR image
|
||||
'''
|
||||
x = imresize_np(x, scale=1 / sf)
|
||||
return x
|
||||
|
||||
|
||||
def srmd_degradation(x, k, sf=3):
|
||||
''' blur + bicubic downsampling
|
||||
Args:
|
||||
x: HxWxC image, [0, 1]
|
||||
k: hxw, double
|
||||
sf: down-scale factor
|
||||
Return:
|
||||
downsampled LR image
|
||||
Reference:
|
||||
@inproceedings{zhang2018learning,
|
||||
title={Learning a single convolutional super-resolution network for multiple degradations},
|
||||
author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
|
||||
booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
|
||||
pages={3262--3271},
|
||||
year={2018}
|
||||
}
|
||||
'''
|
||||
x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
|
||||
x = bicubic_degradation(x, sf=sf)
|
||||
return x
|
||||
|
||||
|
||||
def dpsr_degradation(x, k, sf=3):
|
||||
''' bicubic downsampling + blur
|
||||
Args:
|
||||
x: HxWxC image, [0, 1]
|
||||
k: hxw, double
|
||||
sf: down-scale factor
|
||||
Return:
|
||||
downsampled LR image
|
||||
Reference:
|
||||
@inproceedings{zhang2019deep,
|
||||
title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
|
||||
author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
|
||||
booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
|
||||
pages={1671--1681},
|
||||
year={2019}
|
||||
}
|
||||
'''
|
||||
x = bicubic_degradation(x, sf=sf)
|
||||
x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
|
||||
return x
|
||||
|
||||
|
||||
def classical_degradation(x, k, sf=3):
|
||||
''' blur + downsampling
|
||||
Args:
|
||||
x: HxWxC image, [0, 1]/[0, 255]
|
||||
k: hxw, double
|
||||
sf: down-scale factor
|
||||
Return:
|
||||
downsampled LR image
|
||||
'''
|
||||
x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
|
||||
st = 0
|
||||
return x[st::sf, st::sf, ...]
|
||||
|
||||
|
||||
def add_sharpening(img, weight=0.5, radius=50, threshold=10):
|
||||
"""USM sharpening. borrowed from real-ESRGAN
|
||||
Input image: I; Blurry image: B.
|
||||
1. K = I + weight * (I - B)
|
||||
2. Mask = 1 if abs(I - B) > threshold, else: 0
|
||||
3. Blur mask:
|
||||
4. Out = Mask * K + (1 - Mask) * I
|
||||
Args:
|
||||
img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
|
||||
weight (float): Sharp weight. Default: 1.
|
||||
radius (float): Kernel size of Gaussian blur. Default: 50.
|
||||
threshold (int):
|
||||
"""
|
||||
if radius % 2 == 0:
|
||||
radius += 1
|
||||
blur = cv2.GaussianBlur(img, (radius, radius), 0)
|
||||
residual = img - blur
|
||||
mask = np.abs(residual) * 255 > threshold
|
||||
mask = mask.astype('float32')
|
||||
soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
|
||||
|
||||
K = img + weight * residual
|
||||
K = np.clip(K, 0, 1)
|
||||
return soft_mask * K + (1 - soft_mask) * img
|
||||
|
||||
|
||||
def add_blur_1(img, sf=4):
|
||||
wd2 = 4.0 + sf
|
||||
wd = 2.0 + 0.2 * sf
|
||||
|
||||
wd2 = wd2 / 4
|
||||
wd = wd / 4
|
||||
|
||||
if random.random() < 0.5:
|
||||
l1 = wd2 * random.random()
|
||||
l2 = wd2 * random.random()
|
||||
k = anisotropic_Gaussian(
|
||||
ksize=random.randint(2, 11) + 3,
|
||||
theta=random.random() * np.pi,
|
||||
l1=l1,
|
||||
l2=l2)
|
||||
else:
|
||||
k = fspecial('gaussian',
|
||||
random.randint(2, 4) + 3, wd * random.random())
|
||||
img = ndimage.filters.convolve(
|
||||
img, np.expand_dims(k, axis=2), mode='mirror')
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def add_resize(img, sf=4):
|
||||
rnum = np.random.rand()
|
||||
if rnum > 0.8:
|
||||
sf1 = random.uniform(1, 2)
|
||||
elif rnum < 0.7:
|
||||
sf1 = random.uniform(0.5 / sf, 1)
|
||||
else:
|
||||
sf1 = 1.0
|
||||
img = cv2.resize(
|
||||
img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
img = np.clip(img, 0.0, 1.0)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
|
||||
noise_level = random.randint(noise_level1, noise_level2)
|
||||
rnum = np.random.rand()
|
||||
if rnum > 0.6:
|
||||
img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(
|
||||
np.float32)
|
||||
elif rnum < 0.4:
|
||||
img = img + np.random.normal(0, noise_level / 255.0,
|
||||
(*img.shape[:2], 1)).astype(np.float32)
|
||||
else:
|
||||
L = noise_level2 / 255.
|
||||
D = np.diag(np.random.rand(3))
|
||||
U = orth(np.random.rand(3, 3))
|
||||
conv = np.dot(np.dot(np.transpose(U), D), U)
|
||||
img = img + np.random.multivariate_normal([0, 0, 0], np.abs(
|
||||
L**2 * conv), img.shape[:2]).astype(np.float32)
|
||||
img = np.clip(img, 0.0, 1.0)
|
||||
return img
|
||||
|
||||
|
||||
def add_speckle_noise(img, noise_level1=2, noise_level2=25):
|
||||
noise_level = random.randint(noise_level1, noise_level2)
|
||||
img = np.clip(img, 0.0, 1.0)
|
||||
rnum = random.random()
|
||||
if rnum > 0.6:
|
||||
img += img * np.random.normal(0, noise_level / 255.0,
|
||||
img.shape).astype(np.float32)
|
||||
elif rnum < 0.4:
|
||||
img += img * np.random.normal(0, noise_level / 255.0,
|
||||
(*img.shape[:2], 1)).astype(np.float32)
|
||||
else:
|
||||
L = noise_level2 / 255.
|
||||
D = np.diag(np.random.rand(3))
|
||||
U = orth(np.random.rand(3, 3))
|
||||
conv = np.dot(np.dot(np.transpose(U), D), U)
|
||||
img += img * np.random.multivariate_normal(
|
||||
[0, 0, 0], np.abs(L**2 * conv), img.shape[:2]).astype(np.float32)
|
||||
img = np.clip(img, 0.0, 1.0)
|
||||
return img
|
||||
|
||||
|
||||
def add_Poisson_noise(img):
|
||||
img = np.clip((img * 255.0).round(), 0, 255) / 255.
|
||||
vals = 10**(2 * random.random() + 2.0)
|
||||
if random.random() < 0.5:
|
||||
img = np.random.poisson(img * vals).astype(np.float32) / vals
|
||||
else:
|
||||
img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
|
||||
img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
|
||||
noise_gray = np.random.poisson(img_gray * vals).astype(
|
||||
np.float32) / vals - img_gray
|
||||
img += noise_gray[:, :, np.newaxis]
|
||||
img = np.clip(img, 0.0, 1.0)
|
||||
return img
|
||||
|
||||
|
||||
def add_JPEG_noise(img):
|
||||
quality_factor = random.randint(80, 95)
|
||||
img = cv2.cvtColor(single2uint(img), cv2.COLOR_RGB2BGR)
|
||||
result, encimg = cv2.imencode(
|
||||
'.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
|
||||
img = cv2.imdecode(encimg, 1)
|
||||
img = cv2.cvtColor(uint2single(img), cv2.COLOR_BGR2RGB)
|
||||
return img
|
||||
|
||||
|
||||
def random_crop(lq, hq, sf=4, lq_patchsize=64):
|
||||
h, w = lq.shape[:2]
|
||||
rnd_h = random.randint(0, h - lq_patchsize)
|
||||
rnd_w = random.randint(0, w - lq_patchsize)
|
||||
lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
|
||||
|
||||
rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
|
||||
hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf,
|
||||
rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
|
||||
return lq, hq
|
||||
|
||||
|
||||
def degradation_bsrgan_light(image, sf=4, isp_model=None):
|
||||
"""
|
||||
This is the variant of the degradation model of BSRGAN from the paper
|
||||
"Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
|
||||
sf: scale factor
|
||||
isp_model: camera ISP model
|
||||
Returns
|
||||
img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
|
||||
hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
|
||||
"""
|
||||
image = uint2single(image)
|
||||
_, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
|
||||
|
||||
h1, w1 = image.shape[:2]
|
||||
image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...]
|
||||
h, w = image.shape[:2]
|
||||
|
||||
if sf == 4 and random.random() < scale2_prob:
|
||||
if np.random.rand() < 0.5:
|
||||
image = cv2.resize(
|
||||
image,
|
||||
(int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
else:
|
||||
image = imresize_np(image, 1 / 2, True)
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
sf = 2
|
||||
|
||||
shuffle_order = random.sample(range(7), 7)
|
||||
idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
|
||||
if idx1 > idx2:
|
||||
shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[
|
||||
idx2], shuffle_order[idx1]
|
||||
|
||||
for i in shuffle_order:
|
||||
|
||||
if i == 0:
|
||||
image = add_blur_1(image, sf=sf)
|
||||
elif i == 2:
|
||||
a, b = image.shape[1], image.shape[0]
|
||||
# downsample2
|
||||
if random.random() < 0.8:
|
||||
sf1 = random.uniform(1, 2 * sf)
|
||||
image = cv2.resize(
|
||||
image, (int(1 / sf1 * image.shape[1]),
|
||||
int(1 / sf1 * image.shape[0])),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
else:
|
||||
k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
|
||||
k_shifted = shift_pixel(k, sf)
|
||||
k_shifted = k_shifted / k_shifted.sum()
|
||||
image = ndimage.filters.convolve(
|
||||
image, np.expand_dims(k_shifted, axis=2), mode='mirror')
|
||||
image = image[0::sf, 0::sf, ...]
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
elif i == 3:
|
||||
# downsample3
|
||||
image = cv2.resize(
|
||||
image, (int(1 / sf * a), int(1 / sf * b)),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
elif i == 4:
|
||||
# add Gaussian noise
|
||||
image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2)
|
||||
elif i == 5:
|
||||
# add JPEG noise
|
||||
if random.random() < jpeg_prob:
|
||||
image = add_JPEG_noise(image)
|
||||
|
||||
# add final JPEG compression noise
|
||||
image = add_JPEG_noise(image)
|
||||
image = single2uint(image)
|
||||
return image
|
||||
|
||||
|
||||
def add_blur_2(img, sf=4):
|
||||
wd2 = 4.0 + sf
|
||||
wd = 2.0 + 0.2 * sf
|
||||
if random.random() < 0.5:
|
||||
l1 = wd2 * random.random()
|
||||
l2 = wd2 * random.random()
|
||||
k = anisotropic_Gaussian(
|
||||
ksize=2 * random.randint(2, 11) + 3,
|
||||
theta=random.random() * np.pi,
|
||||
l1=l1,
|
||||
l2=l2)
|
||||
else:
|
||||
k = fspecial('gaussian', 2 * random.randint(2, 11) + 3,
|
||||
wd * random.random())
|
||||
img = ndimage.filters.convolve(
|
||||
img, np.expand_dims(k, axis=2), mode='mirror')
|
||||
return img
|
||||
|
||||
|
||||
def degradation_bsrgan(image, sf=4, isp_model=None):
|
||||
"""
|
||||
This is the variant of the degradation model of BSRGAN from the paper
|
||||
"Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
|
||||
sf: scale factor
|
||||
isp_model: camera ISP model
|
||||
Returns
|
||||
img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
|
||||
hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
|
||||
"""
|
||||
image = uint2single(image)
|
||||
_, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
|
||||
|
||||
h1, w1 = image.shape[:2]
|
||||
image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...]
|
||||
h, w = image.shape[:2]
|
||||
|
||||
if sf == 4 and random.random() < scale2_prob:
|
||||
if np.random.rand() < 0.5:
|
||||
image = cv2.resize(
|
||||
image,
|
||||
(int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
else:
|
||||
image = imresize_np(image, 1 / 2, True)
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
sf = 2
|
||||
|
||||
shuffle_order = random.sample(range(7), 7)
|
||||
idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
|
||||
if idx1 > idx2:
|
||||
shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[
|
||||
idx2], shuffle_order[idx1]
|
||||
|
||||
for i in shuffle_order:
|
||||
|
||||
if i == 0:
|
||||
image = add_blur_2(image, sf=sf)
|
||||
elif i == 1:
|
||||
image = add_blur_2(image, sf=sf)
|
||||
elif i == 2:
|
||||
a, b = image.shape[1], image.shape[0]
|
||||
# downsample2
|
||||
if random.random() < 0.75:
|
||||
sf1 = random.uniform(1, 2 * sf)
|
||||
image = cv2.resize(
|
||||
image, (int(1 / sf1 * image.shape[1]),
|
||||
int(1 / sf1 * image.shape[0])),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
else:
|
||||
k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
|
||||
k_shifted = shift_pixel(k, sf)
|
||||
k_shifted = k_shifted / k_shifted.sum()
|
||||
image = ndimage.filters.convolve(
|
||||
image, np.expand_dims(k_shifted, axis=2), mode='mirror')
|
||||
image = image[0::sf, 0::sf, ...]
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
elif i == 3:
|
||||
# downsample3
|
||||
image = cv2.resize(
|
||||
image, (int(1 / sf * a), int(1 / sf * b)),
|
||||
interpolation=random.choice([1, 2, 3]))
|
||||
image = np.clip(image, 0.0, 1.0)
|
||||
elif i == 4:
|
||||
# add Gaussian noise
|
||||
image = add_Gaussian_noise(image, noise_level1=2, noise_level2=25)
|
||||
elif i == 5:
|
||||
# add JPEG noise
|
||||
if random.random() < jpeg_prob:
|
||||
image = add_JPEG_noise(image)
|
||||
|
||||
# add final JPEG compression noise
|
||||
image = add_JPEG_noise(image)
|
||||
image = single2uint(image)
|
||||
return image
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user