diff --git a/apps/application/workflow/nodes/image_to_video_node/__init__.py b/apps/application/workflow/nodes/image_to_video_node/__init__.py new file mode 100644 index 00000000000..2a6d404093b --- /dev/null +++ b/apps/application/workflow/nodes/image_to_video_node/__init__.py @@ -0,0 +1,2 @@ +# coding=utf-8 +from .image_to_video_node import ImageToVideoNode diff --git a/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py b/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py new file mode 100644 index 00000000000..88ca29cb67e --- /dev/null +++ b/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py @@ -0,0 +1,248 @@ +# coding=utf-8 +import base64 +import uuid_utils.compat as uuid +import requests +from functools import reduce +from typing import List + +from django.db.models import QuerySet +from django.utils.translation import gettext_lazy as _, gettext +from langchain_core.messages import BaseMessage, HumanMessage, AIMessage +from rest_framework import serializers + +from application.workflow.common import WorkflowType +from application.workflow.i_node import INode +from application.workflow.message.struct.content import NodeInfo, Position +from application.workflow.message.struct.text_content import TextContent +from application.workflow.status import Status +from common.utils.common import bytes_to_uploaded_file +from knowledge.models import FileSourceType, File +from oss.serializers.file import FileSerializer, mime_types +from models_provider.tools import get_model_instance_by_model_workspace_id +from common.utils.logger import maxkb_logger + + +class ImageToVideoNodeSerializer(serializers.Serializer): + model_id = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_("Model id")) + model_id_type = serializers.CharField(required=False, default='custom', label=_("Model id type")) + model_id_reference = serializers.ListField(required=False, child=serializers.CharField(), allow_empty=True, + label=_("Reference Field")) + prompt = serializers.CharField(required=True, label=_("Prompt word (positive)")) + negative_prompt = serializers.CharField(required=False, label=_("Prompt word (negative)"), + allow_null=True, allow_blank=True) + dialogue_number = serializers.IntegerField(required=False, default=0, + label=_("Number of multi-round conversations")) + dialogue_type = serializers.CharField(required=False, default='NODE', + label=_("Conversation storage type")) + is_result = serializers.BooleanField(required=False, + label=_('Whether to return content')) + model_params_setting = serializers.JSONField(required=False, default=dict, + label=_("Model parameter settings")) + first_frame_url = serializers.ListField(required=True, label=_("First frame url")) + last_frame_url = serializers.ListField(required=False, label=_("Last frame url")) + + +class ImageToVideoNode(INode): + serializer_class = ImageToVideoNodeSerializer + supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL] + type = 'image-to-video-node' + + def execute(self): + maxkb_logger.info(f'[ImageToVideoNode] execute START, node_id={self.get_node_id()}') + workflow_params = self.get_workflow_parameters() + node_params = self.get_parameters() + + model_id = node_params.get('model_id') + model_id_type = node_params.get('model_id_type', 'custom') + model_id_reference = node_params.get('model_id_reference') + prompt = node_params.get('prompt', '') + negative_prompt = node_params.get('negative_prompt', '') + dialogue_number = node_params.get('dialogue_number', 0) + dialogue_type = node_params.get('dialogue_type', 'NODE') + is_result = node_params.get('is_result', False) + model_params_setting = node_params.get('model_params_setting') + first_frame_url_ref = node_params.get('first_frame_url') + last_frame_url_ref = node_params.get('last_frame_url') + + workflow_type = self.get_workflow_type() + if workflow_type in (WorkflowType.KNOWLEDGE, WorkflowType.TOOL): + history_chat_record = [] + chat_id = None + chat_record_id = None + workspace_id = workflow_params.get('workspace_id') + else: + history_chat_record = workflow_params.get('history_chat_record', []) + chat_id = workflow_params.get('chat_id') + chat_record_id = workflow_params.get('chat_record_id') + workspace_id = workflow_params.get('workspace_id') + + if model_id_type == 'reference' and model_id_reference: + reference_data = self.workflow_manage.get_reference_field( + model_id_reference[0], model_id_reference[1:], + ) + if reference_data and isinstance(reference_data, dict): + model_id = reference_data.get('model_id', model_id) + model_params_setting = reference_data.get('model_params_setting') + + if not model_id: + raise Exception(_('Model is not allowed to be empty')) + + if first_frame_url_ref is None or first_frame_url_ref == []: + raise ValueError(_("First frame url cannot be empty")) + + first_frame_url = self.workflow_manage.get_reference_field( + first_frame_url_ref[0], first_frame_url_ref[1:]) + + last_frame_url = None + if last_frame_url_ref is not None and last_frame_url_ref != []: + last_frame_url = self.workflow_manage.get_reference_field( + last_frame_url_ref[0], last_frame_url_ref[1:]) + + ttv_model = get_model_instance_by_model_workspace_id(model_id, workspace_id, + **(model_params_setting or {})) + + history_message = self._get_history_message(history_chat_record, dialogue_number) + self.write_context('history_message', [ + {'content': message.content, 'role': message.type} + for message in (history_message or []) + ]) + + question = self.workflow_manage.generate_prompt(prompt) + self.write_context('question', question) + + message_list = [*history_message, question] + self.write_context('message_list', message_list) + self.write_context('dialogue_type', dialogue_type) + self.write_context('negative_prompt', self.workflow_manage.generate_prompt(negative_prompt)) + self.write_context('first_frame_url', first_frame_url) + self.write_context('last_frame_url', last_frame_url) + + first_frame_url = self._get_file_base64(first_frame_url) + last_frame_url = self._get_file_base64(last_frame_url) + + video_urls = ttv_model.generate_video(question, negative_prompt, first_frame_url, last_frame_url) + maxkb_logger.info(f'[ImageToVideoNode] generate_video result: {video_urls is not None}, node_id={self.get_node_id()}') + + if video_urls is None or video_urls == '': + raise Exception(gettext('Failed to generate video')) + + file_name = 'generated_video.mp4' + if isinstance(video_urls, str) and video_urls.startswith('http'): + video_urls = requests.get(video_urls).content + + file = bytes_to_uploaded_file(video_urls, file_name) + file_url = self._upload_file(file, workflow_type, workflow_params) + + video_label = f'' + video_list = [{'file_id': file_url.split('/')[-1], 'file_name': file_name, 'url': file_url}] + + self.write_context('answer', video_label) + self.write_context('video', video_list) + self.write_context('chat_model', ttv_model) + + if is_result: + node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.SUCCESS) + self.write(TextContent(str(uuid.uuid7()), video_label, Status.SUCCESS, node_info, Position(self.get_node_id()))) + + def _get_file_base64(self, image_url): + try: + if isinstance(image_url, list): + image_url = image_url[0].get('file_id') if 'file_id' in image_url[0] else image_url[0].get('url') + if isinstance(image_url, str) and not image_url.startswith('http'): + file = QuerySet(File).filter(id=image_url).first() + file_bytes = file.get_bytes() + file_type = file.file_name.split(".")[-1].lower() + content_type = mime_types.get(file_type, 'application/octet-stream') + encoded_bytes = base64.b64encode(file_bytes) + return f'data:{content_type};base64,{encoded_bytes.decode()}' + return image_url + except Exception as e: + raise ValueError(gettext("Failed to obtain the image")) + + def _upload_file(self, file, workflow_type, workflow_params): + if workflow_type == WorkflowType.KNOWLEDGE: + return self._upload_knowledge_file(file, workflow_params) + if workflow_type == WorkflowType.TOOL: + return self._upload_tool_file(file, workflow_params) + return self._upload_application_file(file, workflow_params) + + def _upload_knowledge_file(self, file, workflow_params): + knowledge_id = workflow_params.get('knowledge_id') + meta = { + 'debug': False, + 'knowledge_id': knowledge_id + } + file_url = FileSerializer(data={ + 'file': file, + 'meta': meta, + 'source_id': knowledge_id, + 'source_type': FileSourceType.KNOWLEDGE.value + }).upload() + return file_url + + def _upload_tool_file(self, file, workflow_params): + tool_id = workflow_params.get('tool_id') + meta = { + 'debug': False, + 'tool_id': tool_id, + } + file_url = FileSerializer(data={ + 'file': file, + 'meta': meta, + 'source_id': tool_id, + 'source_type': FileSourceType.TOOL.value + }).upload() + return file_url + + def _upload_application_file(self, file, workflow_params): + application_id = workflow_params.get('application_id') + chat_id = workflow_params.get('chat_id') + debug = workflow_params.get('debug', False) + meta = { + 'debug': debug, + 'chat_id': chat_id, + 'application_id': application_id, + } + file_url = FileSerializer(data={ + 'file': file, + 'meta': meta, + 'source_id': application_id, + 'source_type': FileSourceType.APPLICATION.value + }).upload() + return file_url + + def _generate_history_ai_message(self, chat_record): + for val in chat_record.details.values(): + if self.node.id == val['node_id'] and 'image_list' in val: + if val['dialogue_type'] == 'WORKFLOW': + return chat_record.get_ai_message() + image_list = val['image_list'] + return AIMessage(content=[ + *[{'type': 'image_url', 'image_url': {'url': f'{file_url}'}} for file_url in image_list] + ]) + return chat_record.get_ai_message() + + def _get_history_message(self, history_chat_record, dialogue_number): + start_index = len(history_chat_record) - dialogue_number + history_message = reduce(lambda x, y: [*x, *y], [ + [self._generate_history_human_message(history_chat_record[index]), + self._generate_history_ai_message(history_chat_record[index])] + for index in + range(start_index if start_index > 0 else 0, len(history_chat_record))], []) + return history_message + + def _generate_history_human_message(self, chat_record): + for data in chat_record.details.values(): + if self.node.id == data['node_id'] and 'image_list' in data: + image_list = data['image_list'] + if len(image_list) == 0 or data['dialogue_type'] == 'WORKFLOW': + return HumanMessage(content=chat_record.problem_text) + return HumanMessage(content=data['question']) + return HumanMessage(content=chat_record.problem_text) + + @staticmethod + def reset_message_list(message_list: List[BaseMessage], answer_text): + result = [{'role': 'user' if isinstance(message, HumanMessage) else 'ai', 'content': message.content} for + message in message_list] + result.append({'role': 'ai', 'content': answer_text}) + return result