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:
wenmeng.zwm
2023-08-30 09:46:54 +08:00
200 changed files with 25960 additions and 5717 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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
View File

@@ -124,6 +124,8 @@ replace.sh
result.png
result.jpg
result.mp4
runs/
ckpt/
# Pytorch
*.pth

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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,

View File

@@ -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_modulesrankalpha 需与训练参数相同
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'))

View File

@@ -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 ''

View File

@@ -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(

View File

@@ -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')

View File

@@ -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,

View 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>&nbsp &nbspEnglish
</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`.

View 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>
中文&nbsp &nbsp<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`两列的数据集.

View File

@@ -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)

View File

@@ -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}')

View File

@@ -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,

View 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']

View File

@@ -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

View File

@@ -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,

View File

@@ -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
#

View File

@@ -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",

View File

@@ -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",

View File

@@ -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",

View File

@@ -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",

View File

@@ -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})

View File

@@ -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

View File

@@ -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()

View File

@@ -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

View File

@@ -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()

View File

@@ -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

View File

@@ -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])

View File

@@ -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

View File

@@ -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']
}

View File

@@ -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'

View File

@@ -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(

View File

@@ -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

View 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)

View File

@@ -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)

View File

@@ -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

View File

@@ -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

View File

@@ -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)

View File

@@ -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')

View File

@@ -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

View File

@@ -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)

View 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={},
)

View 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()

View File

@@ -0,0 +1,5 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .autoencoder import *
from .embedder import *
from .unet_i2v import *

View 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

View 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)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,2 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os

View 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
# -----------------------------------------------------------------------------

View 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

View 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

View File

@@ -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}')

View 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

View File

@@ -1,2 +1,3 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .stable_diffusion import StableDiffusion
from .stable_diffusion_xl import StableDiffusionXL

View File

@@ -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:

View File

@@ -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)

View File

@@ -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']

View 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={},
)

View File

@@ -0,0 +1,5 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .autoencoder import *
from .embedder import *
from .unet_v2v import *

View File

@@ -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

View File

@@ -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)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,2 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os

View 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'
# -----------------------------------------------------------------------------

View File

@@ -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()

View File

@@ -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

View 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

View File

@@ -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

View 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

View 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()

View 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={},
)

View File

@@ -0,0 +1 @@
# Copyright (c) Alibaba, Inc. and its affiliates.

View File

@@ -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

View File

@@ -0,0 +1,3 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .palette import *

View File

@@ -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)

View File

@@ -0,0 +1,4 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .pidinet import *
from .sketch_simplification import *

View File

@@ -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

View File

@@ -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

View File

@@ -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

View 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

View 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)

View 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/'

View File

@@ -0,0 +1,5 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from .samplers import *
from .tokenizers import *
from .transforms import *

View 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')

View 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

View 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

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View 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')

View File

@@ -0,0 +1,2 @@
from .clip import *
from .midas import *

View 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)

View 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

View File

@@ -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 *

View 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