开源AIGC学习—文生图模型服务封装

原创

平常心

修改于 2024-03-15 15:14:55

2350

修改于 2024-03-15 15:14:55

文章被收录于专栏：个人总结系列

一、模型本地化

详细内容见：开源AIGC学习—文生图模型本地运行

二、服务封装

当前算法模型采用Python + Flask 方式进行Rest API方式进行服务封装，对应封装代码说明如下：

from gevent import pywsgi
from flask import Flask
from flask_restful import Resource, Api, reqparse
from flask_cors import CORS
import torch
import cv2
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

task = Tasks.text_to_image_synthesis
model_id = '/mnt/d/aigc_model/modelscope/damo/multi-modal_chinese_stable_diffusion_v1'


app = Flask(__name__)
CORS(app, resources={r"/api/*": {"origins": "*"}})
api = Api(app)


task = Tasks.text_to_image_synthesis
model_id = '/mnt/d/aigc_model/modelscope/damo/multi-modal_chinese_stable_diffusion_v1'
pipe = pipeline(task=task, model=model_id)

parser = reqparse.RequestParser()
parser.add_argument('prompts', type=str, help='Inputs for text to image', location='form')

class TextToImag(Resource):
    def post(self):
        args = parser.parse_args()
        text = args['prompts']
        output = pipe({'text': text})
        result = cv2.imwrite('/mnt/d/aigc_result/result.png', output['output_imgs'][0])
        return {'response': result}

api.add_resource(TextToImag, '/api/text2imag')
if __name__ == '__main__':
    server = pywsgi.WSGIServer(('0.0.0.0', 2000), app)
    server.serve_forever()

也可以直接返回图片，但是文生图模型推理耗时比较长，选择异步方案比较好，然后图片生成后在消息知会用户。

# -- utf-8 ---
from gevent import pywsgi
from flask import Flask, send_file
from flask_restful import Resource, Api, reqparse
from flask_cors import CORS
import torch
import cv2
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

task = Tasks.text_to_image_synthesis
model_id = '/mnt/d/aigc_model/modelscope/damo/multi-modal_chinese_stable_diffusion_v1'


app = Flask(__name__)
CORS(app, resources={r"/api/*": {"origins": "*"}})
api = Api(app)


task = Tasks.text_to_image_synthesis
model_id = '/mnt/d/aigc_model/modelscope/damo/multi-modal_chinese_stable_diffusion_v1'
pipe = pipeline(task=task, model=model_id)

parser = reqparse.RequestParser()
parser.add_argument('prompts', type=str, help='Inputs for text to image', location='form')

class TextToImag(Resource):
    def post(self):
        args = parser.parse_args()
        text = args['prompts']
        output = pipe({'text': text})
        output_path = '/mnt/d/aigc_result/result.png'
        cv2.imwrite(output_path, output['output_imgs'][0])
        return send_file(output_path, mimetype='image/jpeg')

api.add_resource(TextToImag, '/api/text2imag')

if __name__ == '__main__':
    server = pywsgi.WSGIServer(('0.0.0.0', 2000), app)
    server.serve_forever()

返回效果展示