\n'
- f'
' # noqa
+ f'
' # noqa
f' 
\n' # noqa
f'
\n' # noqa
- f' \n' # noqa
+ f'
\n' # noqa
f' 
\n' # noqa
f'
\n' # noqa
f' \n' # noqa
f'')
sec_ret_str = f'{label}' if show_label else ''
+ elif type == Media.VIDEO_PAIR:
+ assert isinstance(content, (list, tuple)) and len(content) == 2
+ ret_str = f'\n' # noqa
+ ret_str += f' \n'
+ f' \n' # noqa
+ f' \n' # noqa
+ f' \n' # noqa
+ f'
')
+ sec_ret_str = f'{label}' if show_label else ''
else:
raise NotImplementedError
if cols_span > 1:
diff --git a/scepter/studio/chatbot/chatbot.py b/scepter/studio/chatbot/chatbot.py
index af0377f..7bf29bc 100644
--- a/scepter/studio/chatbot/chatbot.py
+++ b/scepter/studio/chatbot/chatbot.py
@@ -12,15 +12,13 @@ import re
import string
import sys
import threading
+import warnings
import cv2
import gradio as gr
import numpy as np
import torch
import transformers
-from diffusers import CogVideoXImageToVideoPipeline
-from diffusers.utils import export_to_video
-from gradio_imageslider import ImageSlider
from PIL import Image
from transformers import AutoModel, AutoTokenizer
@@ -29,6 +27,7 @@ from scepter.modules.utils.config import Config
from scepter.modules.utils.directory import get_md5
from scepter.modules.utils.file_system import FS
from scepter.studio.utils.env import init_env
+from importlib.metadata import version
from .example import get_examples
from .utils import load_image
@@ -51,6 +50,11 @@ class ChatBotUI(object):
is_debug=False,
language='en',
root_work_dir='./'):
+ try:
+ from diffusers import CogVideoXImageToVideoPipeline
+ from diffusers.utils import export_to_video
+ except Exception as e:
+ print(f"Import diffusers failed, please install or upgrade diffusers. Error information: {e}")
cfg = Config(cfg_file=cfg_general_file)
cfg.WORK_DIR = os.path.join(root_work_dir, cfg.WORK_DIR)
@@ -58,25 +62,27 @@ class ChatBotUI(object):
FS.make_dir(cfg.WORK_DIR)
cfg = init_env(cfg)
self.cache_dir = cfg.WORK_DIR
- self.chatbot_examples = get_examples(self.cache_dir)
+ self.chatbot_examples = get_examples(self.cache_dir) if not cfg.get('SKIP_EXAMPLES', False) else []
self.model_cfg_dir = cfg.MODEL.EDIT_MODEL.MODEL_CFG_DIR
self.model_yamls = glob.glob(os.path.join(self.model_cfg_dir,
'*.yaml'))
self.model_choices = dict()
+ self.default_model_name = ''
for i in self.model_yamls:
- model_name = '.'.join(i.split('/')[-1].split('.')[:-1])
- self.model_choices[model_name] = i
- print('Models: ', self.model_choices)
-
- self.model_name = cfg.MODEL.EDIT_MODEL.DEFAULT
- assert self.model_name in self.model_choices
- model_cfg = Config(load=True,
- cfg_file=self.model_choices[self.model_name])
+ model_cfg = Config(load=True, cfg_file=i)
+ model_name = model_cfg.NAME
+ if model_cfg.IS_DEFAULT: self.default_model_name = model_name
+ self.model_choices[model_name] = model_cfg
+ print('Models: ', self.model_choices.keys())
+ assert len(self.model_choices) > 0
+ if self.default_model_name == "": self.default_model_name = self.model_choices.keys()[0]
+ self.model_name = self.default_model_name
self.pipe = ACEInference()
- self.pipe.init_from_cfg(model_cfg)
+ self.pipe.init_from_cfg(self.model_choices[self.default_model_name])
self.max_msgs = 20
-
self.enable_i2v = cfg.get('ENABLE_I2V', False)
+ self.gradio_version = version('gradio')
+
if self.enable_i2v:
self.i2v_model_dir = cfg.MODEL.I2V.MODEL_DIR
self.i2v_model_name = cfg.MODEL.I2V.MODEL_NAME
@@ -167,6 +173,7 @@ class ChatBotUI(object):
]
def create_ui(self):
+
css = '.chatbot.prose.md {opacity: 1.0 !important} #chatbot {opacity: 1.0 !important}'
with gr.Blocks(css=css,
title='Chatbot',
@@ -177,7 +184,8 @@ class ChatBotUI(object):
self.history_result = gr.State(value={})
self.retry_msg = gr.State(value='')
with gr.Group():
- with gr.Row(equal_height=True):
+ self.ui_mode = gr.State(value='legacy')
+ with gr.Row(equal_height=True, visible=False) as self.chat_group:
with gr.Column(visible=True) as self.chat_page:
self.chatbot = gr.Chatbot(
height=600,
@@ -192,7 +200,7 @@ class ChatBotUI(object):
size='sm')
with gr.Column(visible=False) as self.editor_page:
- with gr.Tabs():
+ with gr.Tabs(visible=False) as self.upload_tabs:
with gr.Tab(id='ImageUploader',
label='Image Uploader',
visible=True) as self.upload_tab:
@@ -201,7 +209,7 @@ class ChatBotUI(object):
interactive=True,
type='pil',
image_mode='RGB',
- sources='upload',
+ sources=['upload'],
elem_id='image_uploader',
format='png')
with gr.Row():
@@ -209,10 +217,9 @@ class ChatBotUI(object):
value='Submit',
elem_id='upload_submit')
self.ext_btn_1 = gr.Button(value='Exit')
-
+ with gr.Tabs(visible=False) as self.edit_tabs:
with gr.Tab(id='ImageEditor',
- label='Image Editor',
- visible=False) as self.edit_tab:
+ label='Image Editor') as self.edit_tab:
self.mask_type = gr.Dropdown(
label='Mask Type',
choices=[
@@ -275,13 +282,23 @@ class ChatBotUI(object):
self.ext_btn_2 = gr.Button(value='Exit')
with gr.Tab(id='ImageViewer',
- label='Image Viewer',
- visible=False) as self.image_view_tab:
- self.image_viewer = ImageSlider(
- label='Image',
- type='pil',
- show_download_button=True,
- elem_id='image_viewer')
+ label='Image Viewer') as self.image_view_tab:
+ if self.gradio_version >= '5.0.0':
+ self.image_viewer = gr.Image(
+ label='Image',
+ type='pil',
+ show_download_button=True,
+ elem_id='image_viewer')
+ else:
+ try:
+ from gradio_imageslider import ImageSlider
+ except Exception as e:
+ print(f"Import gradio_imageslider failed, please install.")
+ self.image_viewer = ImageSlider(
+ label='Image',
+ type='pil',
+ show_download_button=True,
+ elem_id='image_viewer')
self.ext_btn_3 = gr.Button(value='Exit')
@@ -300,11 +317,30 @@ class ChatBotUI(object):
self.ext_btn_4 = gr.Button(value='Exit')
+ with gr.Row(equal_height=True, visible=True) as self.legacy_group:
+ with gr.Column():
+ self.legacy_image_uploader = gr.Image(
+ height=550,
+ interactive=True,
+ type='pil',
+ image_mode='RGB',
+ elem_id='legacy_image_uploader',
+ format='png')
+ with gr.Column():
+ self.legacy_image_viewer = gr.Image(
+ label='Image',
+ height=550,
+ type='pil',
+ interactive=False,
+ show_download_button=True,
+ elem_id='image_viewer')
+
+
with gr.Accordion(label='Setting', open=False):
with gr.Row():
self.model_name_dd = gr.Dropdown(
choices=self.model_choices,
- value=self.model_name,
+ value=self.default_model_name,
label='Model Version')
with gr.Row():
@@ -315,39 +351,63 @@ class ChatBotUI(object):
label='Negative Prompt',
container=False)
+ with gr.Row():
+ # REFINER_PROMPT
+ self.refiner_prompt = gr.Textbox(
+ value=self.pipe.input.get("refiner_prompt", ""),
+ visible=self.pipe.input.get("refiner_prompt", None) is not None,
+ placeholder=
+ 'Prompt used for refiner',
+ label='Refiner Prompt',
+ container=False)
+
+
with gr.Row():
with gr.Column(scale=8, min_width=500):
with gr.Row():
self.step = gr.Slider(minimum=1,
maximum=1000,
- value=20,
+ value=self.pipe.input.get("sample_steps", 20),
+ visible=self.pipe.input.get("sample_steps", None) is not None,
label='Sample Step')
self.cfg_scale = gr.Slider(
minimum=1.0,
maximum=20.0,
- value=4.5,
+ value=self.pipe.input.get("guide_scale", 4.5),
+ visible=self.pipe.input.get("guide_scale", None) is not None,
label='Guidance Scale')
self.rescale = gr.Slider(minimum=0.0,
maximum=1.0,
- value=0.5,
+ value=self.pipe.input.get("guide_rescale", 0.5),
+ visible=self.pipe.input.get("guide_rescale", None) is not None,
label='Rescale')
+ self.refiner_scale = gr.Slider(minimum=-0.1,
+ maximum=1.0,
+ value=self.pipe.input.get("refiner_scale", 0.5),
+ visible=self.pipe.input.get("refiner_scale", None) is not None,
+ label='Refiner Scale')
self.seed = gr.Slider(minimum=-1,
maximum=10000000,
value=-1,
label='Seed')
self.output_height = gr.Slider(
minimum=256,
- maximum=1024,
- value=512,
+ maximum=1440,
+ value=self.pipe.input.get("output_height", 1024),
+ visible=self.pipe.input.get("output_height", None) is not None,
label='Output Height')
self.output_width = gr.Slider(
minimum=256,
- maximum=1024,
- value=512,
+ maximum=1440,
+ value=self.pipe.input.get("output_width", 1024),
+ visible=self.pipe.input.get("output_width", None) is not None,
label='Output Width')
with gr.Column(scale=1, min_width=50):
self.use_history = gr.Checkbox(value=False,
label='Use History')
+ self.use_ace = gr.Checkbox(value=self.pipe.input.get("use_ace", True),
+ visible=self.pipe.input.get("use_ace", None) is not None,
+ label='Use ACE')
self.video_auto = gr.Checkbox(
value=False,
label='Auto Gen Video',
@@ -384,7 +444,7 @@ class ChatBotUI(object):
visible=True)
with gr.Row():
- inst = """
+ self.chatbot_inst = """
**Instruction**:
1. Click 'Upload' button to upload one or more images as input images.
@@ -398,12 +458,25 @@ class ChatBotUI(object):
8. If you find our work valuable, we invite you to refer to the [ACE Page](https://ali-vilab.github.io/ace-page/) for comprehensive information.
"""
- gr.Markdown(value=inst)
+
+ self.legacy_inst = """
+ **Instruction**:
+
+ 1. You can edit the image by uploading it; if no image is uploaded, an image will be generated from text..
+ 2. Enter '@' in the text box will exhibit all images in the gallery.
+ 3. Select the image you wish to edit from the gallery, and its Image ID will be displayed in the text box.
+ 4. **Important** To render text on an image, please ensure to include a space between each letter. For instance, "add text 'g i r l' on the mask area of @xxxxx".
+ 5. To perform multi-step editing, partial editing, inpainting, outpainting, and other operations, please click the Chatbot Checkbox to enable the conversational editing mode and follow the relevant instructions..
+ 6. If you find our work valuable, we invite you to refer to the [ACE Page](https://ali-vilab.github.io/ace-page/) for comprehensive information.
+
+ """
+
+ self.instruction = gr.Markdown(value=self.legacy_inst)
with gr.Row(variant='panel',
equal_height=True,
show_progress=False):
- with gr.Column(scale=1, min_width=100):
+ with gr.Column(scale=1, min_width=100, visible=False) as self.upload_panel:
self.upload_btn = gr.Button(value=upload_sty +
' Upload',
variant='secondary')
@@ -413,12 +486,16 @@ class ChatBotUI(object):
label='Instruction',
container=False)
with gr.Column(scale=1, min_width=100):
- self.chat_btn = gr.Button(value=chat_sty + ' Chat',
+ self.chat_btn = gr.Button(value='Generate',
variant='primary')
with gr.Column(scale=1, min_width=100):
self.retry_btn = gr.Button(value=refresh_sty +
' Retry',
variant='secondary')
+ with gr.Column(scale=1, min_width=100):
+ self.mode_checkbox = gr.Checkbox(
+ value=False,
+ label='ChatBot')
with gr.Column(scale=(1 if self.enable_i2v else 0),
min_width=0):
self.video_gen_btn = gr.Button(value=video_sty +
@@ -453,19 +530,77 @@ class ChatBotUI(object):
lock.acquire()
del self.pipe
torch.cuda.empty_cache()
- model_cfg = Config(load=True,
- cfg_file=self.model_choices[model_name])
self.pipe = ACEInference()
- self.pipe.init_from_cfg(model_cfg)
+ self.pipe.init_from_cfg(self.model_choices[model_name])
self.model_name = model_name
lock.release()
- return model_name, gr.update(), gr.update()
+ return (model_name, gr.update(), gr.update(),
+ gr.Slider(
+ value=self.pipe.input.get("sample_steps", 20),
+ visible=self.pipe.input.get("sample_steps", None) is not None),
+ gr.Slider(
+ value=self.pipe.input.get("guide_scale", 4.5),
+ visible=self.pipe.input.get("guide_scale", None) is not None),
+ gr.Slider(
+ value=self.pipe.input.get("guide_rescale", 0.5),
+ visible=self.pipe.input.get("guide_rescale", None) is not None),
+ gr.Slider(
+ value=self.pipe.input.get("output_height", 1024),
+ visible=self.pipe.input.get("output_height", None) is not None),
+ gr.Slider(
+ value=self.pipe.input.get("output_width", 1024),
+ visible=self.pipe.input.get("output_width", None) is not None),
+ gr.Textbox(
+ value=self.pipe.input.get("refiner_prompt", ""),
+ visible=self.pipe.input.get("refiner_prompt", None) is not None),
+ gr.Slider(
+ value=self.pipe.input.get("refiner_scale", 0.5),
+ visible=self.pipe.input.get("refiner_scale", None) is not None
+ ),
+ gr.Checkbox(
+ value=self.pipe.input.get("use_ace", True),
+ visible=self.pipe.input.get("use_ace", None) is not None
+ )
+ )
self.model_name_dd.change(
change_model,
inputs=[self.model_name_dd],
- outputs=[self.model_name_dd, self.chatbot, self.text])
+ outputs=[
+ self.model_name_dd, self.chatbot, self.text,
+ self.step,
+ self.cfg_scale, self.rescale, self.output_height,
+ self.output_width, self.refiner_prompt, self.refiner_scale,
+ self.use_ace])
+
+
+ def mode_change(mode_check):
+ if mode_check:
+ # ChatBot
+ return (
+ gr.Row(visible=False),
+ gr.Row(visible=True),
+ gr.Button(value='Generate'),
+ gr.State(value='chatbot'),
+ gr.Column(visible=True),
+ gr.Markdown(value=self.chatbot_inst)
+ )
+ else:
+ # Legacy
+ return (
+ gr.Row(visible=True),
+ gr.Row(visible=False),
+ gr.Button(value=chat_sty + ' Chat'),
+ gr.State(value='legacy'),
+ gr.Column(visible=False),
+ gr.Markdown(value=self.legacy_inst)
+ )
+ self.mode_checkbox.change(mode_change, inputs=[self.mode_checkbox],
+ outputs=[self.legacy_group, self.chat_group,
+ self.chat_btn, self.ui_mode,
+ self.upload_panel, self.instruction])
+
########################################
def generate_gallery(text, images):
@@ -522,6 +657,9 @@ class ChatBotUI(object):
fps,
seed,
progress=gr.Progress(track_tqdm=True)):
+
+ from diffusers.utils import export_to_video
+
generator = torch.Generator(device='cuda').manual_seed(seed)
img_ids = re.findall('@(.*?)[ ,;.?$]', message)
if len(img_ids) == 0:
@@ -592,7 +730,11 @@ class ChatBotUI(object):
outputs=[self.history, self.chatbot, self.text, self.gallery])
########################################
- def run_chat(message,
+ def run_chat(
+ message,
+ legacy_image,
+ ui_mode,
+ use_ace,
extend_prompt,
history,
images,
@@ -601,6 +743,8 @@ class ChatBotUI(object):
negative_prompt,
cfg_scale,
rescale,
+ refiner_prompt,
+ refiner_scale,
step,
seed,
output_h,
@@ -612,12 +756,25 @@ class ChatBotUI(object):
video_fps,
video_seed,
progress=gr.Progress(track_tqdm=True)):
+ legacy_img_ids = []
+ if ui_mode == 'legacy':
+ if legacy_image is not None:
+ history, images, img_id = self.add_uploaded_image_to_history(
+ legacy_image, history, images)
+ legacy_img_ids.append(img_id)
retry_msg = message
gen_id = get_md5(message)[:12]
save_path = os.path.join(self.cache_dir, f'{gen_id}.png')
img_ids = re.findall('@(.*?)[ ,;.?$]', message)
history_io = None
+
+ if len(img_ids) < 1:
+ img_ids = legacy_img_ids
+ for img_id in img_ids:
+ if f'@{img_id}' not in message:
+ message = f'@{img_id} ' + message
+
new_message = message
if len(img_ids) > 0:
@@ -676,6 +833,9 @@ class ChatBotUI(object):
guide_scale=cfg_scale,
guide_rescale=rescale,
seed=seed,
+ refiner_prompt=refiner_prompt,
+ refiner_scale=refiner_scale,
+ use_ace=use_ace
)
img = imgs[0]
@@ -784,21 +944,25 @@ class ChatBotUI(object):
while len(history) >= self.max_msgs:
history.pop(0)
- return history, images, history_result, self.get_history(
- history), gr.update(value=''), gr.update(
- visible=False), retry_msg
+ return (history, images, gr.Image(value=save_path),
+ history_result, self.get_history(
+ history), gr.update(), gr.update(
+ visible=False), retry_msg)
chat_inputs = [
+ self.legacy_image_uploader, self.ui_mode, self.use_ace,
self.extend_prompt, self.history, self.images, self.use_history,
self.history_result, self.negative_prompt, self.cfg_scale,
- self.rescale, self.step, self.seed, self.output_height,
+ self.rescale, self.refiner_prompt, self.refiner_scale,
+ self.step, self.seed, self.output_height,
self.output_width, self.video_auto, self.video_step,
self.video_frames, self.video_cfg_scale, self.video_fps,
self.video_seed
]
chat_outputs = [
- self.history, self.images, self.history_result, self.chatbot,
+ self.history, self.images, self.legacy_image_viewer,
+ self.history_result, self.chatbot,
self.text, self.gallery, self.retry_msg
]
@@ -859,6 +1023,8 @@ class ChatBotUI(object):
prompt=[prompt] * img_num,
negative_prompt=[''] * img_num,
seed=seed,
+ refiner_prompt=self.pipe.input.get("refiner_prompt", ""),
+ refiner_scale=self.pipe.input.get("refiner_scale", 0.0),
)
img = imgs[0]
@@ -904,14 +1070,16 @@ class ChatBotUI(object):
return (gr.update(visible=True,
scale=1), gr.update(visible=True, scale=1),
gr.update(visible=True), gr.update(visible=False),
- gr.update(visible=False), gr.update(visible=False))
+ gr.update(visible=False), gr.update(visible=False),
+ gr.update(visible=True))
self.upload_btn.click(upload_image,
inputs=[],
outputs=[
self.chat_page, self.editor_page,
self.upload_tab, self.edit_tab,
- self.image_view_tab, self.video_view_tab
+ self.image_view_tab, self.video_view_tab,
+ self.upload_tabs
])
########################################
@@ -926,13 +1094,19 @@ class ChatBotUI(object):
]
if len(imgs) > 0:
if len(imgs) == 2:
- view_img = copy.deepcopy(imgs)
+ if self.gradio_version >= '5.0.0':
+ view_img = copy.deepcopy(imgs[-1])
+ else:
+ view_img = copy.deepcopy(imgs)
edit_img = copy.deepcopy(imgs[-1])
else:
- view_img = [
- copy.deepcopy(imgs[-1]),
- copy.deepcopy(imgs[-1])
- ]
+ if self.gradio_version >= '5.0.0':
+ view_img = copy.deepcopy(imgs[-1])
+ else:
+ view_img = [
+ copy.deepcopy(imgs[-1]),
+ copy.deepcopy(imgs[-1])
+ ]
edit_img = copy.deepcopy(imgs[-1])
return (gr.update(visible=True,
@@ -941,11 +1115,12 @@ class ChatBotUI(object):
gr.update(visible=False), gr.update(visible=True),
gr.update(visible=True), gr.update(visible=False),
gr.update(value=edit_img),
- gr.update(value=view_img), gr.update(value=None))
+ gr.update(value=view_img), gr.update(value=None),
+ gr.update(visible=True))
else:
return (gr.update(), gr.update(), gr.update(), gr.update(),
gr.update(), gr.update(), gr.update(), gr.update(),
- gr.update())
+ gr.update(), gr.update())
elif isinstance(evt.value, dict) and evt.value.get(
'component', '') == 'video':
value = evt.value['value']['video']['path']
@@ -953,11 +1128,12 @@ class ChatBotUI(object):
scale=1), gr.update(visible=True, scale=1),
gr.update(visible=False), gr.update(visible=False),
gr.update(visible=False), gr.update(visible=True),
- gr.update(), gr.update(), gr.update(value=value))
+ gr.update(), gr.update(), gr.update(value=value),
+ gr.update())
else:
return (gr.update(), gr.update(), gr.update(), gr.update(),
gr.update(), gr.update(), gr.update(), gr.update(),
- gr.update())
+ gr.update(), gr.update())
self.chatbot.select(edit_image,
outputs=[
@@ -965,16 +1141,17 @@ class ChatBotUI(object):
self.upload_tab, self.edit_tab,
self.image_view_tab, self.video_view_tab,
self.image_editor, self.image_viewer,
- self.video_viewer
+ self.video_viewer, self.edit_tabs
])
- self.image_viewer.change(lambda x: x,
- inputs=self.image_viewer,
- outputs=self.image_viewer)
+ if self.gradio_version < '5.0.0':
+ self.image_viewer.change(lambda x: x,
+ inputs=self.image_viewer,
+ outputs=self.image_viewer)
########################################
def submit_upload_image(image, history, images):
- history, images = self.add_uploaded_image_to_history(
+ history, images, _ = self.add_uploaded_image_to_history(
image, history, images)
return gr.update(visible=False), gr.update(
visible=True), gr.update(
@@ -1207,7 +1384,7 @@ class ChatBotUI(object):
history.append(
(None,
f'This is uploaded image:\n {img_str} image ID is: {img_id}'))
- return history, images
+ return history, images, img_id
def run_gr(cfg):
diff --git a/scepter/studio/inference/inference_manager/infer_runer.py b/scepter/studio/inference/inference_manager/infer_runer.py
index 288ae9c..e9f6fc7 100644
--- a/scepter/studio/inference/inference_manager/infer_runer.py
+++ b/scepter/studio/inference/inference_manager/infer_runer.py
@@ -6,6 +6,7 @@ from scepter.modules.inference.largen_inference import LargenInference
from scepter.modules.inference.pixart_inference import PixArtInference
from scepter.modules.inference.sd3_inference import SD3Inference
from scepter.modules.inference.stylebooth_inference import StyleboothInference
+from scepter.modules.inference.cogvideox_inference import CogVideoXInference
from scepter.modules.utils.logger import get_logger
@@ -111,6 +112,8 @@ class PipelineManager():
PipelineBuilder = SD3Inference
elif pipeline_name.startswith('FLUX'):
PipelineBuilder = FluxInference
+ elif pipeline_name.startswith('COGVIDEO'):
+ PipelineBuilder = CogVideoXInference
else:
PipelineBuilder = DiffusionInference
new_inference = PipelineBuilder(logger=self.logger)
diff --git a/scepter/studio/inference/inference_ui/component_names.py b/scepter/studio/inference/inference_ui/component_names.py
index c119167..d0b2aec 100644
--- a/scepter/studio/inference/inference_ui/component_names.py
+++ b/scepter/studio/inference/inference_ui/component_names.py
@@ -98,6 +98,8 @@ class DiffusionUIName():
self.sample = 'Sampler'
self.sample_steps = 'Sample Steps'
self.image_number = 'Images Number'
+ self.num_frames = 'Number of Frames'
+ self.fps = 'Number of FPS'
self.resolutions_height = 'Output Height'
self.resolutions_width = 'Output Width'
self.negative_prompt = 'Negative Prompt'
@@ -121,6 +123,8 @@ class DiffusionUIName():
self.sample = '采样器'
self.sample_steps = '采样步数'
self.image_number = '图片数量'
+ self.num_frames = '视频帧数量'
+ self.fps = '视频帧率'
self.resolutions_height = '输出高度'
self.resolutions_width = '输出宽度'
self.negative_prompt = '负向提示'
diff --git a/scepter/studio/inference/inference_ui/control_ui.py b/scepter/studio/inference/inference_ui/control_ui.py
index 81c7884..4936ab5 100644
--- a/scepter/studio/inference/inference_ui/control_ui.py
+++ b/scepter/studio/inference/inference_ui/control_ui.py
@@ -100,6 +100,7 @@ class ControlUI(UIBase):
label=self.component_names.control_model,
choices=self.controller_choices,
value=self.controller_default,
+ allow_custom_value=True,
interactive=True)
with gr.Column(scale=1, min_width=0):
self.cond_button = gr.Button('Extract')
diff --git a/scepter/studio/inference/inference_ui/diffusion_ui.py b/scepter/studio/inference/inference_ui/diffusion_ui.py
index 691239d..ca78256 100644
--- a/scepter/studio/inference/inference_ui/diffusion_ui.py
+++ b/scepter/studio/inference/inference_ui/diffusion_ui.py
@@ -26,7 +26,7 @@ class DiffusionUI(UIBase):
self.default_resolutions = pipe_manager.pipeline_level_modules[
now_pipeline].paras.RESOLUTIONS
self.default_input = pipe_manager.pipeline_level_modules[
- now_pipeline].input
+ now_pipeline].input_cfg
self.diffusion_paras = self.load_all_paras()
# deal with resolution
@@ -66,7 +66,6 @@ class DiffusionUI(UIBase):
if value is not None and cur_default.get(
key.lower()) not in value:
value.append(cur_default.get(key.lower()))
-
return diffusion_paras
def load_all_paras(self):
@@ -115,22 +114,15 @@ class DiffusionUI(UIBase):
label=self.component_names.resolutions_height,
choices=[key for key in self.cur_h_level_dict.keys()],
value=default_res[0],
+ allow_custom_value=True,
interactive=True)
with gr.Column(scale=1):
self.output_width = gr.Dropdown(
label=self.component_names.resolutions_width,
choices=self.cur_h_level_dict[default_res[0]],
value=default_res[1],
+ allow_custom_value=True,
interactive=True)
- with gr.Row(equal_height=True):
- self.image_number = gr.Slider(
- label=self.component_names.image_number,
- minimum=self.cur_paras.SAMPLES.get('MIN', 1),
- maximum=self.cur_paras.SAMPLES.get('MAX', 4),
- step=1,
- value=self.cur_paras.SAMPLES.get('DEFAULT', 1),
- visible=self.cur_paras.SAMPLES.get('VISIBLE', True),
- interactive=True)
with gr.Row(equal_height=True):
self.sample_steps = gr.Slider(
label=self.component_names.sample_steps,
@@ -139,7 +131,6 @@ class DiffusionUI(UIBase):
step=1,
value=self.cur_paras.SAMPLE_STEPS.get('DEFAULT', 30),
interactive=True)
-
self.guide_scale = gr.Slider(
label=self.component_names.guide_scale,
minimum=self.cur_paras.GUIDE_SCALE.get('MIN', 1),
@@ -156,6 +147,32 @@ class DiffusionUI(UIBase):
value=self.cur_paras.GUIDE_RESCALE.get('DEFAULT', 0.5),
visible=self.cur_paras.GUIDE_RESCALE.get('VISIBLE', True),
interactive=True)
+ with gr.Row(equal_height=True):
+ self.fps = gr.Slider(
+ label=self.component_names.fps,
+ minimum=self.cur_paras.FPS.get('MIN', 1),
+ maximum=self.cur_paras.FPS.get('MAX', 50),
+ step=1,
+ value=self.cur_paras.FPS.get('DEFAULT', 8),
+ visible=self.cur_paras.FPS.get('VISIBLE', True),
+ interactive=True)
+ self.num_frames = gr.Slider(
+ label=self.component_names.num_frames,
+ minimum=self.cur_paras.NUM_FRAMES.get('MIN', 1),
+ maximum=self.cur_paras.NUM_FRAMES.get('MAX', 100),
+ step=1,
+ value=self.cur_paras.NUM_FRAMES.get('DEFAULT', 49),
+ visible=self.cur_paras.NUM_FRAMES.get('VISIBLE', True),
+ interactive=True)
+ with gr.Row(equal_height=True):
+ self.image_number = gr.Slider(
+ label=self.component_names.image_number,
+ minimum=self.cur_paras.SAMPLES.get('MIN', 1),
+ maximum=self.cur_paras.SAMPLES.get('MAX', 4),
+ step=1,
+ value=self.cur_paras.SAMPLES.get('DEFAULT', 1),
+ visible=self.cur_paras.SAMPLES.get('VISIBLE', True),
+ interactive=True)
with gr.Row(equal_height=True):
with gr.Column(scale=1):
self.seed_random = gr.Checkbox(
@@ -177,6 +194,8 @@ class DiffusionUI(UIBase):
'output_height': self.output_height,
'output_width': self.output_width,
'image_number': self.image_number,
+ 'num_frames': self.num_frames,
+ 'fps': self.fps,
'sample_steps': self.sample_steps,
'guide_scale': self.guide_scale,
'guide_rescale': self.guide_rescale,
diff --git a/scepter/studio/inference/inference_ui/gallery_ui.py b/scepter/studio/inference/inference_ui/gallery_ui.py
index 7d2cc6d..571e1f2 100644
--- a/scepter/studio/inference/inference_ui/gallery_ui.py
+++ b/scepter/studio/inference/inference_ui/gallery_ui.py
@@ -1,6 +1,7 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
+import imageio
from collections import OrderedDict
import gradio as gr
@@ -183,8 +184,10 @@ class GalleryUI(UIBase):
'guide_scale':
args.pop('guide_scale'),
'guide_rescale':
- args.pop('guide_rescale')
+ args.pop('guide_rescale'),
}
+ if 'num_frames' in args:
+ pipeline_input.update({'num_frames': args['num_frames']})
args.update({'input': pipeline_input})
def appedix_init(args):
@@ -217,7 +220,7 @@ class GalleryUI(UIBase):
load_init(args)
results = current_pipeline(**args)
- images = []
+ images, videos = [], []
before_images = []
if 'images' in results:
images_tensor = results['images'] * 255
@@ -226,6 +229,11 @@ class GalleryUI(UIBase):
1, 2, 0).cpu().numpy().astype(np.uint8))
for idx in range(images_tensor.shape[0])
]
+ if 'videos' in results:
+ videos = [
+ (video.permute(1, 2, 3, 0).cpu().numpy() * 255).astype(np.uint8)
+ for video in results['videos']
+ ]
if 'before_refine_images' in results and results[
'before_refine_images'] is not None:
before_refine_images_tensor = results['before_refine_images'] * 255
@@ -236,7 +244,7 @@ class GalleryUI(UIBase):
]
if 'seed' in results:
print(results['seed'])
- print(images, before_images)
+ # print(images, before_images)
if args['largen_state']:
largen_history.extend(images)
@@ -250,12 +258,27 @@ class GalleryUI(UIBase):
f'cur_gallery_{i}.jpg')
img.save(save_image)
save_list.append(save_image)
- images = save_list
+ ret_data = save_list
+ else:
+ ret_data = images
+
+ if len(videos) > 0:
+ save_list = []
+ fps = args.get('fps', 8)
+ for i, video in enumerate(videos):
+ save_video_path = os.path.join(self.local_work_dir,
+ f'cur_gallery_{i}.mp4')
+ writer = imageio.get_writer(save_video_path, fps=fps)
+ for frame in video:
+ writer.append_data(np.array(frame))
+ writer.close()
+ save_list.append(save_video_path)
+ ret_data = save_list
return (
gr.Column(visible=len(before_images) > 0),
before_images,
- images,
+ ret_data,
largen_history,
gr.update(value=largen_history),
)
diff --git a/scepter/studio/inference/inference_ui/model_manage_ui.py b/scepter/studio/inference/inference_ui/model_manage_ui.py
index b1e8319..55b9135 100644
--- a/scepter/studio/inference/inference_ui/model_manage_ui.py
+++ b/scepter/studio/inference/inference_ui/model_manage_ui.py
@@ -188,10 +188,11 @@ class ModelManageUI(UIBase):
diffusion_ui.cur_h_level_dict = h_level_dict
default_input = self.pipe_manager.pipeline_level_modules[
- now_pipeline].input
+ now_pipeline].input_cfg
cur_paras = diffusion_ui.get_default(diffusion_ui.diffusion_paras,
default_input)
diffusion_ui.cur_paras = cur_paras
+
return (
diffusion_model,
gr.Dropdown(value=all_module_name['first_stage_model']),
@@ -224,7 +225,11 @@ class ModelManageUI(UIBase):
gr.Slider(value=cur_paras.GUIDE_SCALE.get('DEFAULT', 7.5),
visible=cur_paras.GUIDE_SCALE.get('VISIBLE', True)),
gr.Slider(value=cur_paras.GUIDE_RESCALE.get('DEFAULT', 0.5),
- visible=cur_paras.GUIDE_RESCALE.get('VISIBLE', True)))
+ visible=cur_paras.GUIDE_RESCALE.get('VISIBLE', True)),
+ gr.Slider(value=cur_paras.NUM_FRAMES.get('DEFAULT', 49),
+ visible=cur_paras.NUM_FRAMES.get('VISIBLE', False)),
+ gr.Slider(value=cur_paras.FPS.get('DEFAULT', 8),
+ visible=cur_paras.FPS.get('VISIBLE', False)))
self.diffusion_model.change(
diffusion_model_change,
@@ -240,6 +245,7 @@ class ModelManageUI(UIBase):
diffusion_ui.prompt_prefix, diffusion_ui.output_height,
diffusion_ui.sampler, diffusion_ui.discretization,
diffusion_ui.sample_steps, diffusion_ui.guide_scale,
- diffusion_ui.guide_rescale
+ diffusion_ui.guide_rescale, diffusion_ui.num_frames,
+ diffusion_ui.fps
],
queue=True)
diff --git a/scepter/studio/preprocess/caption_editor_ui/component_names.py b/scepter/studio/preprocess/caption_editor_ui/component_names.py
index 651d72d..2242e3f 100644
--- a/scepter/studio/preprocess/caption_editor_ui/component_names.py
+++ b/scepter/studio/preprocess/caption_editor_ui/component_names.py
@@ -22,7 +22,8 @@ class CreateDatasetUIName():
self.dataset_type = 'Dataset Type'
self.dataset_type_name = {
'scepter_txt2img': 'Text2Image Generation',
- 'scepter_img2img': 'Image Edit Generation'
+ 'scepter_img2img': 'Image Edit Generation',
+ 'scepter_txt2vid': 'Text2Video Generation'
}
self.user_data_name = (
f'Current Dataset Name. Changes of dataset name take '
@@ -34,30 +35,35 @@ class CreateDatasetUIName():
'scepter_txt2img':
f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/3D_example_csv.zip',
'scepter_img2img':
- f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/hed_pair.zip'
+ f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/hed_pair.zip',
+ 'scepter_txt2vid':
+ f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/video_example.zip'
}
self.default_dataset_zip_str = ' and '.join(
[f'[{k}]({v})' for k, v in self.default_dataset_zip.items()])
self.default_dataset_name = {
'scepter_txt2img': '3D_example',
- 'scepter_img2img': 'hed_example'
+ 'scepter_img2img': 'hed_example',
+ 'scepter_txt2vid': 'video_example'
}
-
self.btn_create_datasets_from_file = 'Create Dataset From File'
self.user_direction = (
'### User Guide: \n' +
f'* {self.btn_create_datasets} button is used to create a new dataset '
". Please make sure to modify the dataset's name and version. After creation, "
- 'you can upload images one by one. \n'
+ 'you can upload images or videos one by one. \n'
f'* The "{self.btn_create_datasets_from_file}" button supports creating a new dataset from '
'a file, currently supporting zip files. For zip files, the format should be consistent'
- " with the one used during training, ensuring it contains an 'images/' folder and a '"
- "train.csv' (which will use the image paths in this file); "
- 'The first line is Target:FILE, Prompt, followed by the format of each line: image path, description.'
+ " with the one used during training, ensuring it contains an 'images/' or 'videos/' folder and a '"
+ "train.csv' (which will use the image or video paths in this file); "
+ 'The first line is Target:FILE, Prompt, followed by the format of each line: image path or video path, '
+ 'description.'
'we also surpport the zip of '
- 'one level subfolder of images whose format are in jpg, jpeg, png, webp. '
+ 'one level subfolder of images or videos whose format are in jpg, jpeg, png, mp4, webp. '
f'See the ZIP examples: {self.default_dataset_zip_str}. \n' # noqa
+ 'Addition, txt2vid data also supports batch upload of txt file list, followed by the format of '
+ 'each line: video path#;#video description '
f'* If you have refreshed the page, please click the {self.refresh_list_button} '
'button to ensure all previously created datasets are visible in the dropdown menu.\n'
'* For processing and training with large-scale data(for example more than 10K samples), '
@@ -97,7 +103,8 @@ class CreateDatasetUIName():
self.dataset_type = '数据集类型'
self.dataset_type_name = {
'scepter_txt2img': '文生图数据',
- 'scepter_img2img': '图像编辑（图生图）数据'
+ 'scepter_img2img': '图像编辑（图生图）数据',
+ 'scepter_txt2vid': '文生视频数据'
}
self.user_data_name = f'当前数据集名称，修改后点{self.modify_data_button}生效'
@@ -108,26 +115,30 @@ class CreateDatasetUIName():
'scepter_txt2img':
f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/3D_example_csv.zip',
'scepter_img2img':
- f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/hed_pair.zip'
+ f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/hed_pair.zip',
+ 'scepter_txt2vid':
+ f'{self.default_dataset_repo}repo?Revision=master&FilePath=datasets/video_example.zip'
}
self.default_dataset_zip_str = ' 和 '.join(
[f'[{k}]({v})' for k, v in self.default_dataset_zip.items()])
self.default_dataset_name = {
'scepter_txt2img': '3D_example',
- 'scepter_img2img': 'hed_example'
+ 'scepter_img2img': 'hed_example',
+ 'scepter_txt2vid': 'video_example'
}
self.btn_create_datasets_from_file = '从文件新建'
self.user_direction = (
'### 使用说明 \n' +
f'* {self.btn_create_datasets} 按钮用于从零新建数据集，请注意修改数据集的name和version，'
- '新建完成后可以逐个上传图片。\n' +
+ '新建完成后可以逐个上传图片或视频。\n' +
f'* {self.btn_create_datasets_from_file} 按钮支持从文件中来新建数据集，目前支持zip文件，'
- '需要保证在文件夹外进行打包，并包含images/文件夹和train.csv(会使用该文件中的图片路径)，首行为Target:FILE,Prompt，'
- '其次每行格式为：图片路径,描述；'
- f'同时我们也支持图像文件的zip包，格式在jpg、jpeg、png或webp。数据ZIP样例路径：{self.default_dataset_zip_str}. \n'
+ '需要保证在文件夹外进行打包，并包含 images/ 或 videos/ 文件夹和train.csv(会使用该文件中的图片或视频路径)，首行为Target:FILE,Prompt，'
+ '其次每行格式为：图片 或 视频 路径,描述；'
+ f'同时我们也支持图像或视频文件的zip包，格式在jpg、jpeg、png、mp4或webp。数据ZIP样例路径：{self.default_dataset_zip_str}； \n'
+ '另外，文生视频数据还支持txt文件列表批量上传，文件每行格式为：视频路径#;#视频描述；\n '
+
- f'* 如果刷新了页面，请点击{self.refresh_list_button} 按钮以确保所有以往创建的数据集在下拉框中可见。\n'
+ f'如果刷新了页面，请点击 {self.refresh_list_button} 按钮以确保所有以往创建的数据集在下拉框中可见\n'
'* 对于大规模数据的处理和训练（数据规模大于1万），建议使用命令行形式\n'
'* 请注意观察系统日志的输出以帮助改进操作。 \n')
# Error or Warning
@@ -154,11 +165,13 @@ class DatasetGalleryUIName():
self.illegal_blank_dataset = 'Illgal or blank dataset is not allowed editing.'
self.delete_blank_dataset = 'Blank dataset is not allowed deleting.'
self.upload_image = 'Upload Target Image'
+ self.upload_video = 'Upload Video'
self.upload_src_image = 'Upload Source Image'
self.upload_src_mask = 'Mask Image'
self.upload_image_btn = '\U00002714' # ✔️
self.cancel_upload_btn = '\U00002716' # ✖️
self.image_caption = 'Image Caption'
+ self.video_caption = 'Video Caption'
self.btn_modify = '\U0001F4DD' # 📝
self.btn_delete = '\U0001f5d1' # 🗑️
@@ -196,10 +209,12 @@ class DatasetGalleryUIName():
f'click{self.btn_reset_edit} to reset edited data，'
f'click{self.btn_cancel_edit} to out of editing mode.')
self.preprocess_choices = [
- 'Image Preprocess', 'Caption Preprocess'
+ 'Image Preprocess', 'Caption Preprocess', 'Caption translation'
]
+ self.preprocess_choices_video = ['Video caption generation', 'Caption translation']
self.preview_target_image = 'Preview Target Image'
+ self.preview_target_video = 'Preview Target Video'
self.preview_src_image = 'Preview Source Image'
self.preview_src_mask_image = 'Preview Source Image Mask'
self.preview_caption = 'Preview Caption'
@@ -211,7 +226,7 @@ class DatasetGalleryUIName():
self.caption_preprocess_btn = 'apply'
self.caption_preview_btn = 'preview'
self.caption_update_mode = 'Caption Update Mode'
- self.caption_update_choices = ['Append', 'Replace']
+ self.caption_update_choices = ['Replace', 'Append']
self.used_device = 'Used Device'
self.used_memory = 'Used Memory'
@@ -234,11 +249,13 @@ class DatasetGalleryUIName():
self.illegal_blank_dataset = '不合法或空白数据集不允许编辑。'
self.delete_blank_dataset = '空白数据集不允许删除。'
self.upload_image = '上传目标图片'
+ self.upload_video = '上传视频'
self.upload_src_image = '上传待编辑图片'
self.upload_src_mask = '蒙版区域'
self.upload_image_btn = '\U00002714' # ✔️
self.cancel_upload_btn = '\U00002716' # ✖️
self.image_caption = '图片描述'
+ self.video_caption = '视频描述'
self.btn_modify = '\U0001F4DD' # 📝
self.dataset_images = f'图片数据，点击{self.btn_modify}进入编辑模式'
@@ -256,8 +273,8 @@ class DatasetGalleryUIName():
self.edit_caption = '编辑描述'
self.batch_caption_generate = '处理范围'
- self.ori_dataset = '原始数据 高({}) * 宽({}) 图像格式({})'
- self.edit_dataset = '可编辑数据 高({}) * 宽({}) 图像格式({})'
+ self.ori_dataset = '原始数据 高({}) * 宽({}) 格式({})'
+ self.edit_dataset = '可编辑数据 高({}) * 宽({}) 格式({})'
self.upload_image_info = '图像信息 高({}) * 宽({})'
self.upload_src_image_info = '源图像信息 高({}) * 宽({})'
@@ -276,8 +293,10 @@ class DatasetGalleryUIName():
f'点击{self.btn_cancel_edit}取消编辑，'
f'点击{self.btn_reset_edit}重置数据，'
f'修改编辑范围可以批量编辑不同范围的数据。')
- self.preprocess_choices = ['图像预处理', '描述生成']
+ self.preprocess_choices = ['图像预处理', '描述生成', '描述翻译']
+ self.preprocess_choices_video = ['视频描述生成', '描述翻译']
self.preview_target_image = '预览图片'
+ self.preview_target_video = '预览视频'
self.preview_src_image = '预览原图'
self.preview_src_mask_image = '预览蒙版'
self.preview_caption = '预览描述'
@@ -288,7 +307,7 @@ class DatasetGalleryUIName():
self.caption_preprocess_btn = '应用'
self.caption_preview_btn = '预览'
self.caption_update_mode = '描述更新方式'
- self.caption_update_choices = ['追加', '替换']
+ self.caption_update_choices = ['替换', '追加']
self.used_device = '使用设备'
self.used_memory = '使用内存'
self.caption_language = '描述语言'
@@ -386,3 +405,33 @@ class Image2ImageDataCardName():
self.illegal_data_err7 = '上传图像失败{}'
self.delete_err1 = '删除失败，数据已经为空了'
self.export_zip_err1 = '压缩文件失败!'
+
+
+class Text2VideoDataCardName():
+ def __init__(self, language='en'):
+ if language == 'en':
+ self.illegal_data_err1 = (
+ 'The list supports only "," or "#;#" as delimiters. '
+ 'The two columns represent video path and description, '
+ 'respectively.')
+ self.illegal_data_err2 = 'Illegal file format'
+ self.illegal_data_err3 = 'File decompression failed, failed to upload to storage!'
+ self.illegal_data_err4 = 'Illegal width({}),height({})'
+ self.illegal_data_err5 = (
+ 'The path should not contain "{}". '
+ 'It should be an OSS path (oss://) or the prefix '
+ 'can be omitted (xxx/xxx)."')
+ self.illegal_data_err6 = 'Video download failed {}'
+ self.illegal_data_err7 = 'Video upload failed {}'
+ self.delete_err1 = 'Deletion failed, the data is already empty.'
+ self.export_zip_err1 = 'Failed to compress the file!'
+ elif language == 'zh':
+ self.illegal_data_err1 = '列表只支持,或#;#作为分割符，两列分别为视频路径/描述'
+ self.illegal_data_err2 = '非法的文件格式'
+ self.illegal_data_err3 = '文件解压失败，上传存储器失败！'
+ self.illegal_data_err4 = '不合法的width({}),height({})'
+ self.illegal_data_err5 = '路径不支持{}，应该为oss路径（oss://）或者省略前缀（xxx/xxx）'
+ self.illegal_data_err6 = '下载视频失败{}'
+ self.illegal_data_err7 = '上传视频失败{}'
+ self.delete_err1 = '删除失败，数据已经为空了'
+ self.export_zip_err1 = '压缩文件失败!'
diff --git a/scepter/studio/preprocess/caption_editor_ui/create_dataset_ui.py b/scepter/studio/preprocess/caption_editor_ui/create_dataset_ui.py
index 26674a4..0b51320 100644
--- a/scepter/studio/preprocess/caption_editor_ui/create_dataset_ui.py
+++ b/scepter/studio/preprocess/caption_editor_ui/create_dataset_ui.py
@@ -15,6 +15,8 @@ from scepter.studio.preprocess.utils.img2img_data_card import \
Image2ImageDataCard
from scepter.studio.preprocess.utils.txt2img_data_card import \
Text2ImageDataCard
+from scepter.studio.preprocess.utils.txt2vid_data_card import \
+ Text2VideoDataCard
from scepter.studio.utils.uibase import UIBase
from tqdm import tqdm
@@ -44,7 +46,9 @@ class CreateDatasetUI(UIBase):
'scepter_txt2img':
Text2ImageDataCard,
'scepter_img2img':
- Image2ImageDataCard
+ Image2ImageDataCard,
+ 'scepter_txt2vid':
+ Text2VideoDataCard
})
self.components_name = CreateDatasetUIName(language)
self.default_dataset_type = list(self.dataset_type_dict.keys())[0]
@@ -466,8 +470,7 @@ class CreateDatasetUI(UIBase):
], [
self.panel_state, self.dataset_name, self.user_dataset_name,
self.sys_log
- ],
- queue=False)
+ ], queue=False)
def show_edit_panel(panel_state, data_name):
if panel_state:
@@ -599,10 +602,9 @@ class CreateDatasetUI(UIBase):
trans_dataset_type, [])
else:
dataset_list = []
- return gr.Dropdown(
- value=dataset_list[-1] if len(dataset_list) > 0 else '',
- choices=dataset_list), self.components_name.system_log.format(
- '')
+ return (gr.Dropdown(
+ value=dataset_list[-1] if len(dataset_list) > 0 else '', choices=dataset_list),
+ self.components_name.system_log.format(''))
manager.user_name.change(dataset_type_change,
inputs=[self.dataset_type, manager.user_name],
diff --git a/scepter/studio/preprocess/caption_editor_ui/dataset_gallery_ui.py b/scepter/studio/preprocess/caption_editor_ui/dataset_gallery_ui.py
index 1fefd8a..6c621a8 100644
--- a/scepter/studio/preprocess/caption_editor_ui/dataset_gallery_ui.py
+++ b/scepter/studio/preprocess/caption_editor_ui/dataset_gallery_ui.py
@@ -5,6 +5,7 @@ from __future__ import annotations
import os.path
import time
+import cv2
import gradio as gr
import imagehash
# from gradio.processing_utils import encode_pil_to_base64
@@ -34,7 +35,14 @@ class DatasetGalleryUI(UIBase):
self.processors_manager = ProcessorsManager(cfg.PROCESSORS,
language=language)
+ self.video_processors_manager = ProcessorsManager(cfg.VIDEO_PROCESSORS,
+ language=language)
+
+ self.trans_processors_manager = ProcessorsManager(cfg.TRANSLATION_PROCESSORS,
+ language=language)
+
self.component_names = DatasetGalleryUIName(language)
+ self.translation_name = self.component_names.preprocess_choices[2]
if create_ins is not None:
self.default_dataset = create_ins.default_dataset
current_info = self.default_dataset.current_record
@@ -236,11 +244,16 @@ class DatasetGalleryUI(UIBase):
with gr.Column(scale=1, min_width=0):
self.btn_reset_edit = gr.Button(
value=self.component_names.btn_reset_edit)
- with gr.Row():
+ with gr.Row() as self.mode_select_edit:
self.preprocess_checkbox = gr.CheckboxGroup(
show_label=False,
choices=self.component_names.preprocess_choices,
value=None)
+ with gr.Row(visible=False) as self.mode_select_video:
+ self.upload_preprocess_video = gr.CheckboxGroup(
+ show_label=False,
+ choices=self.component_names.preprocess_choices_video,
+ value=None)
with gr.Column(variant='panel',
visible=False,
scale=1,
@@ -269,7 +282,8 @@ class DatasetGalleryUI(UIBase):
type='pil',
height=240,
interactive=False)
- with gr.Row():
+ # upload image
+ with gr.Row() as self.upload_image_row:
with gr.Column(scale=1):
self.upload_image = gr.Image(
label=self.component_names.upload_image,
@@ -282,6 +296,23 @@ class DatasetGalleryUI(UIBase):
placeholder='',
value='',
lines=8)
+ # upload video
+ with gr.Row():
+ with gr.Row(variant='panel',
+ visible=self.default_dataset_type ==
+ 'scepter_img2video') as self.upload_src_video:
+ with gr.Column(scale=1):
+ self.upload_video = gr.Video(
+ label=self.component_names.upload_video,
+ sources=['upload']
+ )
+ with gr.Column(scale=1):
+ self.upload_caption = gr.Textbox(
+ label=self.component_names.video_caption,
+ autoscroll=True,
+ placeholder='',
+ value='',
+ lines=8)
with gr.Row():
with gr.Column(min_width=0):
self.upload_button = gr.Button(
@@ -289,7 +320,8 @@ class DatasetGalleryUI(UIBase):
with gr.Column(min_width=0):
self.cancel_button = gr.Button(
value=self.component_names.cancel_upload_btn)
- with gr.Row():
+ # mode_select
+ with gr.Row() as self.mode_select_add:
self.upload_preprocess_checkbox = gr.CheckboxGroup(
show_label=False,
choices=self.component_names.preprocess_choices,
@@ -321,7 +353,7 @@ class DatasetGalleryUI(UIBase):
interactive=True)
self.preview_src_mask_image_tool = gr.State(
value='sketch')
- with gr.Column(scale=1):
+ with gr.Column(scale=1) as self.preview_target_image_panel:
self.preview_taget_image = gr.ImageMask(
label=self.component_names.preview_target_image,
sources=[],
@@ -331,6 +363,12 @@ class DatasetGalleryUI(UIBase):
interactive=True)
self.preview_taget_image_tool = gr.State(
value='sketch')
+ with gr.Column(scale=1, visible=False) as self.preview_target_video_panel:
+ self.preview_target_video = gr.Video(
+ label=self.component_names.preview_target_video,
+ sources=['upload'],
+ interactive=True)
+ self.preview_target_video_tool = gr.State(value='sketch')
with gr.Column(scale=1):
self.preview_caption = gr.Textbox(
label=self.component_names.preview_caption,
@@ -431,7 +469,7 @@ class DatasetGalleryUI(UIBase):
'caption')
default_processor_ins = self.processors_manager.get_processor(
'caption', default_processor_method)
- with gr.Column(scale=1, min_width=0):
+ with gr.Column(scale=1, min_width=0) as self.caption_language_panel:
self.caption_language = gr.Dropdown(
label=self.component_names.
caption_language,
@@ -450,7 +488,7 @@ class DatasetGalleryUI(UIBase):
if len(self.component_names.
caption_update_choices) > 0 else
None)
- with gr.Row():
+ with gr.Row() as self.default_use_local_panel:
default_use_local = default_processor_ins.get_para_by_language(
default_processor_ins.get_language_default
).get('USE_LOCAL', False)
@@ -461,7 +499,7 @@ class DatasetGalleryUI(UIBase):
interactive=True)
with gr.Accordion(
label=self.component_names.advance_setting,
- open=False):
+ open=False) as self.advance_setting_panel:
with gr.Row():
self.sys_prompt = gr.Text(
label=self.component_names.
@@ -933,11 +971,28 @@ class DatasetGalleryUI(UIBase):
queue=False)
def view_mode():
- return gr.Text(value='view')
+ return (gr.Text(value='view'),
+ gr.Column(visible=False),
+ gr.Column(visible=True),
+ gr.Dropdown(choices=self.processors_manager.
+ get_choices('caption'),
+ value=self.processors_manager.
+ get_default('caption')),
+ gr.Column(visible=True),
+ gr.Row(visible=True),
+ gr.Accordion(visible=True),
+ gr.CheckboxGroup(value=None))
self.btn_cancel_edit.click(view_mode,
inputs=[],
- outputs=[self.mode_state],
+ outputs=[self.mode_state,
+ self.preview_target_video_panel,
+ self.preview_target_image_panel,
+ self.caption_preprocess_method,
+ self.caption_language_panel,
+ self.default_use_local_panel,
+ self.advance_setting_panel,
+ self.upload_preprocess_video],
queue=False)
# reset edit information to clean the status of editing
@@ -1056,12 +1111,31 @@ class DatasetGalleryUI(UIBase):
return (gr.Gallery(), gr.Gallery(), gr.Image(), gr.Textbox(),
gr.Text(), gr.Text(), gr.Text(),
self.component_names.system_log.format('None'),
- gr.Text())
+ gr.Text(), gr.Column(visible=False),
+ gr.Column(visible=True),
+ gr.Dropdown(choices=self.processors_manager.
+ get_choices('caption'),
+ value=self.processors_manager.
+ get_default('caption')),
+ gr.Column(visible=True),
+ gr.Row(visible=True),
+ gr.Accordion(visible=True),
+ gr.CheckboxGroup())
is_flg, msg = dataset_ins.apply_changes()
if not is_flg:
return (gr.Gallery(), gr.Gallery(), gr.Image(), gr.Textbox(),
gr.Text(), gr.Text(), gr.Text(),
- self.component_names.system_log.format(msg), gr.Text())
+ self.component_names.system_log.format(msg), gr.Text(),
+ gr.Column(visible=False),
+ gr.Column(visible=True),
+ gr.Dropdown(choices=self.processors_manager.
+ get_choices('caption'),
+ value=self.processors_manager.
+ get_default('caption')),
+ gr.Column(visible=True),
+ gr.Row(visible=True),
+ gr.Accordion(visible=True),
+ gr.Checkboxgroup())
image_list = [
os.path.join(dataset_ins.local_work_dir, v['relative_path'])
for v in dataset_ins.data
@@ -1102,7 +1176,16 @@ class DatasetGalleryUI(UIBase):
ret_src_image_gl, ret_mask,
gr.Textbox(value=current_record.get('caption', '')),
image_info, self.component_names.system_log.format(''),
- gr.Text(value='view'))
+ gr.Text(value='view'), gr.Column(visible=False),
+ gr.Column(visible=True),
+ gr.Dropdown(choices=self.processors_manager.
+ get_choices('caption'),
+ value=self.processors_manager.
+ get_default('caption')),
+ gr.Column(visible=True),
+ gr.Row(visible=True),
+ gr.Accordion(visible=True),
+ gr.CheckboxGroup(value=None))
self.btn_confirm_edit.click(
confirm_edit,
@@ -1110,15 +1193,19 @@ class DatasetGalleryUI(UIBase):
outputs=[
self.gl_dataset_images, self.src_gl_dataset_images,
self.src_mask, self.ori_caption, self.image_info, self.sys_log,
- self.mode_state
+ self.mode_state, self.preview_target_video_panel,
+ self.preview_target_image_panel,
+ self.caption_preprocess_method,
+ self.caption_language_panel,
+ self.default_use_local_panel,
+ self.advance_setting_panel,
+ self.upload_preprocess_video
],
queue=False)
def preprocess_box_change(preprocess_checkbox, dataset_name,
- dataset_type, preview_src_image_tool,
- preview_src_mask_image_tool,
- preview_taget_image_tool):
- image_proc_status, caption_proc_status = False, False
+ dataset_type):
+ image_proc_status, caption_proc_status, translation_proc_status = False, False, False
reverse_status = {
v: id
for id, v in enumerate(self.component_names.preprocess_choices)
@@ -1129,7 +1216,9 @@ class DatasetGalleryUI(UIBase):
image_proc_status = True
elif hit_status == 1:
caption_proc_status = True
- if image_proc_status or caption_proc_status:
+ elif hit_status == 2:
+ translation_proc_status = True
+ if image_proc_status or caption_proc_status or translation_proc_status:
dataset_type = create_dataset.get_trans_dataset_type(
dataset_type)
dataset_ins = create_dataset.dataset_dict.get(
@@ -1143,27 +1232,11 @@ class DatasetGalleryUI(UIBase):
src_image_path = os.path.join(
dataset_ins.local_work_dir,
one_data['edit_src_relative_path'])
- # if preview_src_image_tool == 'sketch':
- # image = Image.open(src_image_path)
- # w, h = image.size
- # ret_src_image = gr.Image(value={
- # 'image': encode_pil_to_base64(image),
- # 'mask': default_mask(w, h)
- # }, visible=True)
- # else:
ret_src_image = gr.Image(value=src_image_path,
visible=True)
src_mask_path = os.path.join(
dataset_ins.local_work_dir,
one_data['edit_src_mask_relative_path'])
- # if preview_src_mask_image_tool == 'sketch':
- # image = Image.open(src_mask_path)
- # w, h = image.size
- # ret_src_mask = gr.Image(value={
- # 'image': encode_pil_to_base64(image),
- # 'mask': default_mask(w, h)
- # }, visible=True)
- # else:
ret_src_mask = gr.Image(value=src_mask_path, visible=True)
ret_src_panel = gr.Column(visible=True)
ret_src_mask_panel = gr.Column(visible=True)
@@ -1174,17 +1247,34 @@ class DatasetGalleryUI(UIBase):
ret_src_mask_panel = gr.Column(visible=False)
image_path = os.path.join(dataset_ins.local_work_dir,
one_data['edit_relative_path'])
- # if preview_taget_image_tool == 'sketch':
- # image = Image.open(image_path)
- # w, h = image.size
- # ret_target_image = gr.Image(value={
- # 'image': encode_pil_to_base64(image),
- # 'mask': default_mask(w, h)
- # })
- # else:
ret_target_image = gr.Image(value=image_path)
ret_caption = gr.Textbox(value=one_data['edit_caption'])
ret_panel = gr.Row(visible=True)
+
+ if translation_proc_status and not caption_proc_status:
+ ret_preprocess_method = gr.Dropdown(choices=self.trans_processors_manager.
+ get_choices('caption'),
+ value=self.trans_processors_manager.get_default
+ ('caption'))
+ ret_use_device = gr.Text(value=self.trans_processors_manager.
+ get_default_device('caption'))
+ ret_use_memory = gr.Text(value=self.trans_processors_manager.
+ get_default_memory('caption'))
+ ret_caption_language = gr.Column(visible=False)
+ ret_use_local = gr.Row(visible=False)
+ ret_advance_setting = gr.Accordion(visible=False)
+ else:
+ ret_caption_language = gr.Column(visible=True)
+ ret_preprocess_method = gr.Dropdown(choices=self.processors_manager.
+ get_choices('caption'),
+ value=self.processors_manager.get_default
+ ('caption'))
+ ret_use_device = gr.Text(value=self.processors_manager.
+ get_default_device('caption'))
+ ret_use_memory = gr.Text(value=self.processors_manager.
+ get_default_memory('caption'))
+ ret_use_local = gr.Row(visible=True)
+ ret_advance_setting = gr.Accordion(visible=True)
else:
ret_src_image = gr.Image()
ret_src_mask = gr.Image()
@@ -1193,22 +1283,29 @@ class DatasetGalleryUI(UIBase):
ret_panel = gr.Row(visible=False)
ret_src_panel = gr.Column(visible=False)
ret_src_mask_panel = gr.Column(visible=False)
+ ret_caption_language = gr.Column()
+ ret_preprocess_method = gr.Dropdown()
+ ret_use_device = gr.Text()
+ ret_use_memory = gr.Text()
+ ret_use_local = gr.Row()
+ ret_advance_setting = gr.Accordion()
ret_image_preprocess_method = gr.Dropdown(
visible=image_proc_status)
return (gr.Column(
- visible=image_proc_status or caption_proc_status),
+ visible=image_proc_status or caption_proc_status or translation_proc_status),
gr.Row(visible=image_proc_status),
- gr.Row(visible=caption_proc_status), ret_panel,
+ gr.Row(visible=caption_proc_status or translation_proc_status), ret_panel,
ret_src_image, ret_src_panel, ret_src_mask,
ret_src_mask_panel, ret_target_image, ret_caption,
+ ret_caption_language, ret_preprocess_method,
+ ret_use_device, ret_use_memory, ret_use_local, ret_advance_setting,
ret_image_preprocess_method)
self.preprocess_checkbox.change(
preprocess_box_change,
inputs=[
self.preprocess_checkbox, create_dataset.dataset_name,
- create_dataset.dataset_type, self.preview_src_image_tool,
- self.preview_src_mask_image_tool, self.preview_taget_image_tool
+ create_dataset.dataset_type
],
outputs=[
self.preprocess_panel,
@@ -1228,6 +1325,15 @@ class DatasetGalleryUI(UIBase):
self.preview_taget_image,
# gr.TextBox()
self.preview_caption,
+ # gr.Column()
+ self.caption_language_panel,
+ self.caption_preprocess_method,
+ self.caption_use_device,
+ self.caption_use_memory,
+ # gr.Row()
+ self.default_use_local_panel,
+ # gr.Accordion()
+ self.advance_setting_panel,
self.image_preprocess_method
],
queue=False)
@@ -1236,8 +1342,7 @@ class DatasetGalleryUI(UIBase):
preprocess_box_change,
inputs=[
self.preprocess_checkbox, create_dataset.dataset_name,
- create_dataset.dataset_type, self.preview_src_image_tool,
- self.preview_src_mask_image_tool, self.preview_taget_image_tool
+ create_dataset.dataset_type
],
outputs=[
self.preprocess_panel,
@@ -1257,10 +1362,143 @@ class DatasetGalleryUI(UIBase):
self.preview_taget_image,
# gr.TextBox()
self.preview_caption,
+ # gr.Column()
+ self.caption_language_panel,
+ self.caption_preprocess_method,
+ self.caption_use_device,
+ self.caption_use_memory,
+ # gr.Row()
+ self.default_use_local_panel,
+ # gr.Accordion()
+ self.advance_setting_panel,
self.image_preprocess_method
],
)
+ def preprocess_box_change_video(preprocess_checkbox, dataset_name, dataset_type):
+ translation_proc_status, caption_proc_status = False, False
+ reverse_status = {
+ v: id
+ for id, v in enumerate(self.component_names.preprocess_choices_video)
+ }
+ for value in preprocess_checkbox:
+ hit_status = reverse_status[value]
+ if hit_status == 0:
+ caption_proc_status = True
+ elif hit_status == 1:
+ translation_proc_status = True
+ if translation_proc_status or caption_proc_status:
+ dataset_type = create_dataset.get_trans_dataset_type(
+ dataset_type)
+ dataset_ins = create_dataset.dataset_dict.get(
+ dataset_type, {}).get(dataset_name, None)
+ edit_index_list = dataset_ins.edit_list
+ if len(edit_index_list) > 0:
+ one_data = dataset_ins.data[edit_index_list[0]]
+ else:
+ one_data = {}
+
+ ret_src_image = gr.Image(visible=False)
+ ret_src_mask = gr.Image(visible=False)
+ ret_src_panel = gr.Column(visible=False)
+ ret_src_mask_panel = gr.Column(visible=False)
+
+ video_path = os.path.join(dataset_ins.local_work_dir,
+ one_data['relative_path'])
+ ret_target_video = gr.Video(value=video_path)
+ ret_target_video_panel = gr.Column(visible=True)
+ ret_target_image = gr.Column(visible=False)
+ ret_caption_language = gr.Column(visible=False)
+ ret_use_local = gr.Row(visible=False)
+ ret_advance_setting = gr.Accordion(visible=False)
+ ret_caption = gr.Textbox(value=one_data['caption'])
+ ret_panel = gr.Row(visible=True)
+ if translation_proc_status:
+ ret_preprocess_method = gr.Dropdown(choices=self.trans_processors_manager.
+ get_choices('caption'),
+ value=self.trans_processors_manager.get_default
+ ('caption'))
+ ret_use_device = gr.Text(value=self.trans_processors_manager.
+ get_default_device('caption'))
+ ret_use_memory = gr.Text(value=self.trans_processors_manager.
+ get_default_memory('caption'))
+ elif caption_proc_status:
+ ret_preprocess_method = gr.Dropdown(choices=self.video_processors_manager.
+ get_choices('caption'),
+ value=self.video_processors_manager.get_default
+ ('caption'))
+ ret_use_device = gr.Text(value=self.video_processors_manager.
+ get_default_device('caption'))
+ ret_use_memory = gr.Text(value=self.video_processors_manager.
+ get_default_memory('caption'))
+ else:
+ ret_src_image = gr.Image()
+ ret_src_mask = gr.Image()
+ ret_target_video = gr.Video()
+ ret_target_video_panel = gr.Column()
+ ret_target_image = gr.Column()
+ ret_caption = gr.Textbox()
+ ret_panel = gr.Row(visible=False)
+ ret_caption_language = gr.Column()
+ ret_use_local = gr.Row()
+ ret_advance_setting = gr.Accordion()
+ ret_src_panel = gr.Column(visible=False)
+ ret_preprocess_method = gr.Dropdown()
+ ret_src_mask_panel = gr.Column(visible=False)
+ ret_use_device = gr.Text()
+ ret_use_memory = gr.Text()
+
+ ret_image_preprocess_method = gr.Dropdown(
+ visible=caption_proc_status)
+ return (gr.Column(
+ visible=caption_proc_status or translation_proc_status),
+ gr.Row(visible=False),
+ gr.Row(visible=caption_proc_status or translation_proc_status),
+ ret_panel, ret_src_image, ret_src_panel, ret_src_mask,
+ ret_src_mask_panel, ret_target_video, ret_target_video_panel,
+ ret_target_image, ret_caption, ret_caption_language,
+ ret_use_local, ret_advance_setting, ret_preprocess_method,
+ ret_use_device, ret_use_memory, ret_image_preprocess_method)
+
+ self.upload_preprocess_video.change(
+ preprocess_box_change_video,
+ inputs=[self.upload_preprocess_video,
+ create_dataset.dataset_name,
+ create_dataset.dataset_type],
+ outputs=[self.preprocess_panel,
+ self.image_preprocess_panel,
+ self.caption_preprocess_panel,
+ # gr.Row()
+ self.preview_panel,
+ # gr.Image()
+ self.preview_src_image,
+ # gr.Column()
+ self.preview_src_panel,
+ # gr.Image()
+ self.preview_src_mask_image,
+ # gr.Column()
+ self.preview_src_mask_panel,
+ # gr.Video()
+ self.preview_target_video,
+ # gr.Column()
+ self.preview_target_video_panel,
+ # gr.Column()
+ self.preview_target_image_panel,
+ # gr.TextBox()
+ self.preview_caption,
+ # gr.Column()
+ self.caption_language_panel,
+ # gr.Row()
+ self.default_use_local_panel,
+ # gr.Accordion()
+ self.advance_setting_panel,
+ self.caption_preprocess_method,
+ self.caption_use_device,
+ self.caption_use_memory,
+ self.image_preprocess_method
+ ]
+ )
+
def upload_preprocess_box_change(preprocess_checkbox, dataset_type,
upload_src_image, upload_src_mask,
upload_image, upload_caption):
@@ -1749,24 +1987,28 @@ class DatasetGalleryUI(UIBase):
],
queue=False)
- def caption_preprocess_method_change(caption_preprocess_method):
- processor_ins = self.processors_manager.get_processor(
- 'caption', caption_preprocess_method)
+ def caption_preprocess_method_change(caption_preprocess_method, dataset_type):
+ dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
+ if dataset_type == 'scepter_txt2vid':
+ processor_ins = self.video_processors_manager.get_processor(
+ 'caption', caption_preprocess_method)
+ else:
+ processor_ins = self.processors_manager.get_processor(
+ 'caption', caption_preprocess_method)
if processor_ins is None:
- return (gr.Text(), gr.Text(), gr.Dropdown(),
- self.component_names.system_log.format(
- 'Load processor failed, processor is None.'))
+ return (gr.Text(),
+ gr.Text(),
+ gr.Dropdown())
language_choice = processor_ins.get_language_choice
language_default = processor_ins.get_language_default
return (gr.Text(value=processor_ins.use_device),
gr.Text(value=f'{processor_ins.use_memory}M'),
gr.Dropdown(choices=language_choice,
- value=language_default),
- self.component_names.system_log.format(''))
+ value=language_default))
self.caption_preprocess_method.change(
caption_preprocess_method_change,
- inputs=[self.caption_preprocess_method],
+ inputs=[self.caption_preprocess_method, create_dataset.dataset_type],
outputs=[
self.caption_use_device, self.caption_use_memory,
self.caption_language
@@ -1774,9 +2016,15 @@ class DatasetGalleryUI(UIBase):
queue=False)
def caption_language_change(caption_language,
- caption_preprocess_method):
- processor_ins = self.processors_manager.get_processor(
- 'caption', caption_preprocess_method)
+ caption_preprocess_method,
+ dataset_type):
+ dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
+ if dataset_type == 'scepter_txt2vid':
+ processor_ins = self.video_processors_manager.get_processor(
+ 'caption', caption_preprocess_method)
+ else:
+ processor_ins = self.processors_manager.get_processor(
+ 'caption', caption_preprocess_method)
para = processor_ins.get_para_by_language(caption_language)
system_prompt = para.get('PROMPT', '')
ret_system_prompt = gr.Text(value=system_prompt,
@@ -1824,7 +2072,7 @@ class DatasetGalleryUI(UIBase):
self.caption_language.change(
caption_language_change,
- inputs=[self.caption_language, self.caption_preprocess_method],
+ inputs=[self.caption_language, self.caption_preprocess_method, create_dataset.dataset_type],
outputs=[
self.sys_prompt, self.max_new_tokens, self.min_new_tokens,
self.num_beams, self.repetition_penalty, self.temperature
@@ -1836,18 +2084,25 @@ class DatasetGalleryUI(UIBase):
repetition_penalty, temperature, use_local,
caption_update_mode, upload_image,
upload_src_image, upload_src_mask,
- upload_caption, dataset_type, dataset_name):
-
+ upload_caption, dataset_type, dataset_name,
+ preprocess_name, preprocess_name_txt2vid):
reverse_update_mode = {
v: idx
for idx, v in enumerate(
self.component_names.caption_update_choices)
}
-
+ dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
+ trans_status = get_trans_status(dataset_type, preprocess_name, preprocess_name_txt2vid)
update_mode = reverse_update_mode.get(caption_update_mode, -1)
-
- processor_ins = self.processors_manager.get_processor(
- 'caption', preprocess_method)
+ if dataset_type == 'scepter_txt2vid' and not trans_status:
+ processor_ins = self.video_processors_manager.get_processor(
+ 'caption', preprocess_method)
+ elif trans_status:
+ processor_ins = self.trans_processors_manager.get_processor(
+ 'caption', preprocess_method)
+ else:
+ processor_ins = self.processors_manager.get_processor(
+ 'caption', preprocess_method)
if processor_ins is None:
sys_log = 'Current processor is illegal'
return gr.Textbox(), gr.Textbox(
@@ -1858,58 +2113,70 @@ class DatasetGalleryUI(UIBase):
sys_log = f'Load processor failed: {msg}'
return gr.Textbox(), gr.Textbox(
), self.component_names.system_log.format(sys_log)
- dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
dataset_ins = create_dataset.dataset_dict.get(
dataset_type, {}).get(dataset_name, None)
if dataset_ins is None:
return gr.Textbox(), gr.Textbox(
), self.component_names.system_log.format('None')
+
if mode_state == 'edit':
edit_index_list = dataset_ins.edit_list
for index in edit_index_list:
one_data = dataset_ins.data[index]
- relative_image_path = one_data.get(
- 'edit_relative_path', one_data['relative_path'])
- target_image = Image.open(
- os.path.join(dataset_ins.meta['local_work_dir'],
- relative_image_path))
-
- if dataset_type == 'scepter_img2img':
+ if (dataset_type == 'scepter_txt2vid' and
+ not trans_status):
+ video_path = one_data.get('video_path', None)
+ kwargs = {
+ 'video_path': video_path
+ }
+ elif trans_status:
+ caption = one_data.get('caption', None)
+ kwargs = {
+ 'caption': caption
+ }
+ else:
relative_image_path = one_data.get(
- 'edit_relative_path',
- one_data['src_relative_path'])
- src_image = Image.open(
+ 'edit_relative_path', one_data['relative_path'])
+ target_image = Image.open(
os.path.join(dataset_ins.meta['local_work_dir'],
relative_image_path))
- relative_src_mask_path = one_data.get(
- 'edit_relative_path',
- one_data['src_mask_relative_path'])
- src_mask_image = Image.open(
- os.path.join(dataset_ins.meta['local_work_dir'],
- relative_src_mask_path))
- else:
- src_image = None
- src_mask_image = None
- kwargs = {
- 'src_image': src_image,
- 'src_mask': src_mask_image,
- 'target_image': target_image,
- 'caption': one_data.get('edit_caption', ''),
- 'preview_src_image': None,
- 'preview_src_mask': None,
- 'preview_target_image': None,
- 'preview_caption': None,
- 'use_preview': False,
- 'use_local': use_local,
- 'sys_prompt': sys_prompt,
- 'max_new_tokens': max_new_tokens,
- 'min_new_tokens': min_new_tokens,
- 'num_beams': num_beams,
- 'repetition_penalty': repetition_penalty,
- 'temperature': temperature,
- 'cache': self.cache
- }
+ if dataset_type == 'scepter_img2img':
+ relative_image_path = one_data.get(
+ 'edit_relative_path',
+ one_data['src_relative_path'])
+ src_image = Image.open(
+ os.path.join(dataset_ins.meta['local_work_dir'],
+ relative_image_path))
+ relative_src_mask_path = one_data.get(
+ 'edit_relative_path',
+ one_data['src_mask_relative_path'])
+ src_mask_image = Image.open(
+ os.path.join(dataset_ins.meta['local_work_dir'],
+ relative_src_mask_path))
+ else:
+ src_image = None
+ src_mask_image = None
+
+ kwargs = {
+ 'src_image': src_image,
+ 'src_mask': src_mask_image,
+ 'target_image': target_image,
+ 'caption': one_data.get('edit_caption', ''),
+ 'preview_src_image': None,
+ 'preview_src_mask': None,
+ 'preview_target_image': None,
+ 'preview_caption': None,
+ 'use_preview': False,
+ 'use_local': use_local,
+ 'sys_prompt': sys_prompt,
+ 'max_new_tokens': max_new_tokens,
+ 'min_new_tokens': min_new_tokens,
+ 'num_beams': num_beams,
+ 'repetition_penalty': repetition_penalty,
+ 'temperature': temperature,
+ 'cache': self.cache
+ }
response = processor_ins(**kwargs)
if update_mode == 0:
@@ -1999,7 +2266,8 @@ class DatasetGalleryUI(UIBase):
self.use_local, self.caption_update_mode, self.upload_image,
self.upload_src_image, self.upload_src_mask,
self.upload_caption, create_dataset.dataset_type,
- create_dataset.dataset_name
+ create_dataset.dataset_name, self.preprocess_checkbox,
+ self.upload_preprocess_video
],
outputs=[self.edit_caption, self.upload_caption, self.sys_log],
queue=False)
@@ -2070,7 +2338,8 @@ class DatasetGalleryUI(UIBase):
preprocess_method, dataset_type, sys_prompt, max_new_tokens,
min_new_tokens, num_beams, repetition_penalty, temperature,
use_local, caption_update_mode, preview_src_image,
- preview_src_mask_image, preview_target_image, preview_caption):
+ preview_src_mask_image, preview_target_image, preview_caption,
+ video_path, preprocess_name, preprocess_name_txt2vid):
reverse_update_mode = {
v: idx
for idx, v in enumerate(
@@ -2078,60 +2347,68 @@ class DatasetGalleryUI(UIBase):
}
update_mode = reverse_update_mode.get(caption_update_mode, -1)
+ dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
+ trans_status = get_trans_status(dataset_type, preprocess_name, preprocess_name_txt2vid)
+ if dataset_type == 'scepter_txt2vid' and not trans_status:
+ processor_ins = self.video_processors_manager.get_processor(
+ 'caption', preprocess_method)
+ kwargs = {
+ 'video_path': video_path
+ }
+ elif trans_status:
+ processor_ins = self.trans_processors_manager.get_processor(
+ 'caption', preprocess_method)
+ kwargs = {
+ 'caption': preview_caption
+ }
+ else:
+ processor_ins = self.processors_manager.get_processor(
+ 'caption', preprocess_method)
+ if isinstance(preview_target_image, dict):
+ prev_target_image = preview_target_image['background']
+ else:
+ prev_target_image = preview_target_image
- processor_ins = self.processors_manager.get_processor(
- 'caption', preprocess_method)
- if processor_ins is None:
- sys_log = 'Current processor is illegal'
- return gr.Textbox(), gr.Textbox(
- ), self.component_names.system_log.format(sys_log)
+ if dataset_type == 'scepter_img2img':
+ if isinstance(preview_src_image, dict):
+ prev_src_image = preview_src_image['background']
+ else:
+ prev_src_image = preview_src_image
+
+ if isinstance(preview_src_mask_image, dict):
+ prev_src_mask = preview_src_mask_image['layers'][0]
+ else:
+ prev_src_mask = preview_src_mask_image
+
+ else:
+ prev_src_image = None
+ prev_src_mask = None
+
+ kwargs = {
+ 'src_image': None,
+ 'src_mask': None,
+ 'target_image': None,
+ 'caption': None,
+ 'preview_src_image': prev_src_image,
+ 'preview_src_mask': prev_src_mask,
+ 'preview_target_image': prev_target_image,
+ 'preview_caption': preview_caption,
+ 'use_preview': True,
+ 'use_local': use_local,
+ 'sys_prompt': sys_prompt,
+ 'max_new_tokens': max_new_tokens,
+ 'min_new_tokens': min_new_tokens,
+ 'num_beams': num_beams,
+ 'repetition_penalty': repetition_penalty,
+ 'temperature': temperature,
+ 'cache': self.cache
+ }
is_flag, msg = processor_ins.load_model()
if not is_flag:
sys_log = f'Load processor failed: {msg}'
return gr.Textbox(), self.component_names.system_log.format(
sys_log)
- dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
-
- if isinstance(preview_target_image, dict):
- prev_target_image = preview_target_image['background']
- else:
- prev_target_image = preview_target_image
-
- if dataset_type == 'scepter_img2img':
- if isinstance(preview_src_image, dict):
- prev_src_image = preview_src_image['background']
- else:
- prev_src_image = preview_src_image
-
- if isinstance(preview_src_mask_image, dict):
- prev_src_mask = preview_src_mask_image['layers'][0]
- else:
- prev_src_mask = preview_src_mask_image
-
- else:
- prev_src_image = None
- prev_src_mask = None
-
- kwargs = {
- 'src_image': None,
- 'src_mask': None,
- 'target_image': None,
- 'caption': None,
- 'preview_src_image': prev_src_image,
- 'preview_src_mask': prev_src_mask,
- 'preview_target_image': prev_target_image,
- 'preview_caption': preview_caption,
- 'use_preview': True,
- 'use_local': use_local,
- 'sys_prompt': sys_prompt,
- 'max_new_tokens': max_new_tokens,
- 'min_new_tokens': min_new_tokens,
- 'num_beams': num_beams,
- 'repetition_penalty': repetition_penalty,
- 'temperature': temperature,
- 'cache': self.cache
- }
response = processor_ins(**kwargs)
if update_mode == 0:
@@ -2156,7 +2433,9 @@ class DatasetGalleryUI(UIBase):
self.num_beams, self.repetition_penalty, self.temperature,
self.use_local, self.caption_update_mode,
self.preview_src_image, self.preview_src_mask_image,
- self.preview_taget_image, self.preview_caption
+ self.preview_taget_image, self.preview_caption,
+ self.preview_target_video, self.preprocess_checkbox,
+ self.upload_preprocess_video
],
outputs=[self.preview_caption, self.sys_log])
@@ -2169,9 +2448,10 @@ class DatasetGalleryUI(UIBase):
if dataset_ins is None or len(dataset_ins) < 1:
return (gr.Gallery(), gr.Gallery(),
gr.Image(), gr.Column(), gr.Row(visible=True),
- gr.Column(), gr.Column(), gr.Image(), gr.Row(), '',
- gr.Markdown(), gr.Textbox(), gr.CheckboxGroup(),
- gr.Column(), gr.Row(), gr.Row(), gr.Column(),
+ gr.Column(), gr.Column(), gr.Image(), gr.Row(),
+ gr.Row(), '', gr.Markdown(), gr.Textbox(),
+ gr.CheckboxGroup(), gr.Column(), gr.Row(),
+ gr.Row(), gr.Column(),
gr.Gallery(), gr.Textbox(), gr.Gallery(),
gr.Image(), gr.Column(), gr.Dropdown(),
gr.Dropdown(value=[]),
@@ -2229,7 +2509,12 @@ class DatasetGalleryUI(UIBase):
ret_edit_mask = gr.Image(visible=False)
return (gr.Gallery(), gr.Gallery(), gr.Image(), gr.Column(),
gr.Row(visible=True), gr.Column(visible=True),
- gr.Column(visible=False), gr.Image(), gr.Row(), '',
+ gr.Column(visible=False), gr.Image(), gr.Row(), gr.Row(),
+ gr.Row(), gr.Row(),
+ gr.Row(visible=False) if dataset_type
+ == 'scepter_txt2vid' else gr.Row(visible=True),
+ gr.Row(visible=True) if dataset_type
+ == 'scepter_txt2vid' else gr.Row(visible=False), '',
gr.Markdown(visible=False), gr.Textbox(),
gr.CheckboxGroup(value=None), gr.Column(visible=False),
gr.Row(visible=False), gr.Row(visible=False),
@@ -2246,7 +2531,17 @@ class DatasetGalleryUI(UIBase):
gr.Row(visible=False), gr.Column(visible=False),
gr.Column(visible=True), gr.Image(),
gr.Row(visible=True) if dataset_type
- == 'scepter_img2img' else gr.Column(visible=False), '',
+ == 'scepter_img2img' else gr.Column(visible=False),
+ gr.Row(visible=True) if dataset_type
+ == 'scepter_txt2vid' else gr.Column(visible=False),
+ gr.Row(visible=False) if dataset_type
+ == 'scepter_txt2vid' else gr.Column(visible=True),
+ gr.Row(visible=False) if dataset_type
+ == 'scepter_txt2vid' else gr.Row(visible=True),
+ gr.Row(visible=False) if dataset_type
+ == 'scepter_txt2vid' else gr.Row(visible=True),
+ gr.Row(visible=True) if dataset_type
+ == 'scepter_txt2vid' else gr.Row(visible=False), '',
gr.Markdown(visible=True) if dataset_type
== 'scepter_img2img' else gr.Markdown(visible=False),
gr.Textbox(value=''), gr.CheckboxGroup(),
@@ -2263,11 +2558,12 @@ class DatasetGalleryUI(UIBase):
if dataset_ins is None:
return (gr.Gallery(), gr.Gallery(),
gr.Image(), gr.Column(), gr.Row(visible=True),
- gr.Column(), gr.Column(), gr.Image(), gr.Row(), '',
+ gr.Column(), gr.Column(), gr.Image(), gr.Row(),
+ gr.Row(), gr.Row(), gr.Row(), gr.Row(), gr.Row(), '',
gr.Markdown(), gr.Textbox(), gr.CheckboxGroup(),
- gr.Column(), gr.Row(), gr.Row(), gr.Column(),
- gr.Gallery(), gr.Textbox(), gr.Gallery(),
- gr.Image(), gr.Column(), gr.Dropdown(),
+ gr.Column(), gr.Row(),
+ gr.Row(), gr.Column(), gr.Gallery(), gr.Textbox(),
+ gr.Gallery(), gr.Image(), gr.Column(), gr.Dropdown(),
gr.Dropdown(value=[]),
self.component_names.system_log.format(
self.component_names.illegal_blank_dataset))
@@ -2328,7 +2624,8 @@ class DatasetGalleryUI(UIBase):
return (ret_gl_gallery, ret_src_gallery, ret_mask,
ret_src_panel, gr.Row(visible=False),
gr.Column(visible=False), gr.Column(visible=False),
- gr.Image(), gr.Row(), '', gr.Markdown(visible=False),
+ gr.Image(), gr.Row(), gr.Row(), gr.Row(), gr.Row(),
+ gr.Row(), gr.Row(), '', gr.Markdown(visible=False),
gr.Textbox(value=''), gr.CheckboxGroup(value=None),
gr.Column(visible=False), gr.Row(visible=False),
gr.Row(visible=False), gr.Column(visible=False),
@@ -2362,6 +2659,16 @@ class DatasetGalleryUI(UIBase):
self.upload_image,
# gr.Row
self.upload_src_image_panel,
+ # gr.Row
+ self.upload_src_video,
+ # gr.Row
+ self.upload_image_row,
+ # gr.Row
+ self.mode_select_add,
+ # gr.Row
+ self.mode_select_edit,
+ # gr.Row
+ self.mode_select_video,
# gr.Markdown
self.upload_image_info,
# gr.Markdown
@@ -2698,6 +3005,22 @@ class DatasetGalleryUI(UIBase):
outputs=[self.upload_image_info],
queue=False)
+ def video_upload(upload_video):
+ if isinstance(upload_video, dict):
+ video = upload_video['video']
+ else:
+ video = upload_video
+ cap = cv2.VideoCapture(video)
+ w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ cap.release()
+ return self.component_names.upload_image_info.format(h, w)
+
+ self.upload_video.upload(video_upload,
+ inputs=[self.upload_video],
+ outputs=[self.upload_image_info],
+ queue=False)
+
def image_src_upload(upload_image):
if isinstance(upload_image, dict):
image = upload_image['image']
@@ -2744,7 +3067,7 @@ class DatasetGalleryUI(UIBase):
)
def add_file(dataset_type, dataset_name, upload_image,
- upload_src_image, upload_src_mask, caption):
+ upload_src_image, upload_src_mask, upload_video, caption):
dataset_type = create_dataset.get_trans_dataset_type(dataset_type)
dataset_ins = create_dataset.dataset_dict.get(
dataset_type, {}).get(dataset_name, None)
@@ -2770,7 +3093,7 @@ class DatasetGalleryUI(UIBase):
if dataset_type == 'scepter_img2img' and src_image is not None and image is None:
w, h = src_image.size
image = Image.new('RGB', (w, h), (0, 0, 0))
- if image is None:
+ if image is None and dataset_type != 'scepter_txt2vid':
return (gr.Image(value=None), gr.Image(value=None),
gr.Image(value=None), gr.Text(value=''), '', '',
gr.Text(value='view'))
@@ -2783,10 +3106,15 @@ class DatasetGalleryUI(UIBase):
return (gr.Image(value=None), gr.Image(value=None),
gr.Image(value=None), gr.Text(value=''), '', '',
gr.Text(value='view'))
- dataset_ins.add_record(image,
- caption,
- src_image=src_image,
- src_mask=src_mask)
+
+ if dataset_type == 'scepter_txt2vid':
+ dataset_ins.add_record(upload_video,
+ caption)
+ else:
+ dataset_ins.add_record(image,
+ caption,
+ src_image=src_image,
+ src_mask=src_mask)
return (gr.Image(value=None), gr.Image(value=None),
gr.Image(value=None), gr.Text(value=''), '', '',
gr.Text(value='view'))
@@ -2796,7 +3124,8 @@ class DatasetGalleryUI(UIBase):
create_dataset.dataset_type,
create_dataset.dataset_name,
self.upload_image, self.upload_src_image,
- self.upload_src_mask, self.upload_caption
+ self.upload_src_mask, self.upload_video,
+ self.upload_caption
],
outputs=[
self.upload_image, self.upload_src_image,
@@ -2830,3 +3159,16 @@ class DatasetGalleryUI(UIBase):
self.edit_caption
],
queue=False)
+
+ def get_trans_status(dataset_type, preprocess_name, preprocess_name_txt2vid):
+ if dataset_type == 'scepter_txt2vid':
+ choices = self.component_names.preprocess_choices_video
+ operation = preprocess_name_txt2vid[0]
+ hit_index = 1
+ else:
+ choices = self.component_names.preprocess_choices
+ operation = preprocess_name[0]
+ hit_index = 2
+
+ reverse_status = {v: id for id, v in enumerate(choices)}
+ return reverse_status.get(operation, -1) == hit_index
diff --git a/scepter/studio/preprocess/processors/caption_processors.py b/scepter/studio/preprocess/processors/caption_processors.py
index b63f9e0..bd719e8 100644
--- a/scepter/studio/preprocess/processors/caption_processors.py
+++ b/scepter/studio/preprocess/processors/caption_processors.py
@@ -15,8 +15,11 @@ from scepter.modules.utils.file_system import FS
import numpy as np
from scepter.studio.preprocess.processors.base_processor import \
BaseCaptionProcessor
+import io
+from decord import cpu, VideoReader, bridge
-__all__ = ['BlipImageBase', 'QWVL', 'QWVLQuantize', 'InternVL15']
+__all__ = ['BlipImageBase', 'QWVL', 'QWVLQuantize', 'InternVL15', 'CogVLM2Llama3Caption',
+ 'OpusMtZhEn', 'OpusMtEnZh']
def get_region(image, mask, mask_id):
locs = np.where(np.array(mask) == mask_id)
@@ -467,3 +470,306 @@ class InternVL15(QWVL):
response = self.model_info['model'].chat(self.model_info['tokenizer'], image, prompt, generation_config=generation_config)
response = response.replace("\n", "").strip()
return response
+
+
+class CogVLM2Llama3Caption(BaseCaptionProcessor):
+ def __init__(self, cfg, language='en'):
+ super().__init__(cfg, language=language)
+ self.model_path = cfg.MODEL_PATH
+ self.model_info = {
+ 'device': 'offline',
+ 'model': None,
+ 'tokenizer': None
+ }
+ self.prompt = cfg.PROMPT
+ self.temperature = cfg.TEMPERATURE
+ self.max_new_tokens = cfg.MAX_NEW_TOKENS
+ self.pad_token_id = cfg.PAD_TOKEN_ID
+ self.top_k = cfg.TOP_K
+ self.top_p = cfg.TOP_P
+ self.TORCH_TYPE = torch.bfloat16 if (torch.cuda.is_available() and
+ torch.cuda.get_device_capability()
+ [0] >= 8) else torch.float16
+
+ def load_model(self):
+ is_flg, msg = super().load_model()
+ if not is_flg:
+ return is_flg, msg
+ if self.model_info['device'] == 'offline':
+ model = None
+ try:
+ from transformers import AutoModelForCausalLM, AutoTokenizer
+ local_model_dir = FS.get_dir_to_local_dir(self.model_path)
+ tokenizer = AutoTokenizer.from_pretrained(
+ local_model_dir,
+ trust_remote_code=True,
+ )
+ model = AutoModelForCausalLM.from_pretrained(
+ local_model_dir,
+ device_map='auto',
+ torch_dtype=self.TORCH_TYPE,
+ trust_remote_code=True
+ ).eval().to(we.device_id)
+ except Exception as e:
+ if model is not None:
+ del model
+ return False, f"Load model error '{e}'"
+ self.model_info['device'] = model.device
+ self.model_info['model'] = model
+ self.model_info['tokenizer'] = tokenizer
+ elif self.model_info['device'] == 'cpu':
+ try:
+ self.model_info['model'].to(we.device_id)
+ self.model_info['device'] = we.device_id
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ except Exception as e:
+ del self.model_info['model']
+ self.model_info['model'] = None
+ self.model_info['device'] = 'offline'
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return False, f"Load model error '{e}'"
+ return True, ''
+
+ def unload_model(self):
+ super().unload_model()
+ if self.delete_instance:
+ self.model_info['device'] = 'offline'
+ if self.model_info['model'] is not None:
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ del self.model_info['model']
+ self.model_info['model'] = None
+ elif (isinstance(self.model_info['device'], numbers.Number)
+ or str(self.model_info['device']).startswith('cuda')):
+ self.model_info['device'] = 'cpu'
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return True, ''
+
+ def load_video(self, video_data, strategy='chat'):
+ bridge.set_bridge('torch')
+ mp4_stream = video_data
+ num_frames = 24
+ decord_vr = VideoReader(io.BytesIO(mp4_stream), ctx=cpu(0))
+
+ frame_id_list = None
+ total_frames = len(decord_vr)
+ if strategy == 'base':
+ clip_end_sec = 60
+ clip_start_sec = 0
+ start_frame = int(clip_start_sec * decord_vr.get_avg_fps())
+ end_frame = min(total_frames,
+ int(clip_end_sec * decord_vr.get_avg_fps())) if clip_end_sec is not None else total_frames
+ frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int)
+ elif strategy == 'chat':
+ timestamps = decord_vr.get_frame_timestamp(np.arange(total_frames))
+ timestamps = [i[0] for i in timestamps]
+ max_second = round(max(timestamps)) + 1
+ frame_id_list = []
+ for second in range(max_second):
+ closest_num = min(timestamps, key=lambda x: abs(x - second))
+ index = timestamps.index(closest_num)
+ frame_id_list.append(index)
+ if len(frame_id_list) >= num_frames:
+ break
+ video_data = decord_vr.get_batch(frame_id_list)
+ video_data = video_data.permute(3, 0, 1, 2)
+ return video_data
+
+ def get_caption(self, prompt, video_data, temperature):
+ strategy = 'chat'
+ video = self.load_video(video_data, strategy=strategy)
+
+ history = []
+ query = prompt
+ model = self.model_info['model']
+ tokenizer = self.model_info['tokenizer']
+ inputs = model.build_conversation_input_ids(
+ tokenizer=tokenizer,
+ query=query,
+ images=[video],
+ history=history,
+ template_version=strategy
+ )
+ inputs = {
+ 'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'),
+ 'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(we.device_id),
+ 'attention_mask': inputs['attention_mask'].unsqueeze(0).to(we.device_id),
+ 'images': [[inputs['images'][0].to(we.device_id).to(self.TORCH_TYPE)]],
+ }
+ gen_kwargs = {
+ "max_new_tokens": self.max_new_tokens,
+ "pad_token_id": self.pad_token_id,
+ "top_k": self.top_k,
+ "do_sample": True,
+ "top_p": self.top_p,
+ "temperature": temperature,
+ }
+ with torch.no_grad():
+ outputs = model.generate(**inputs, **gen_kwargs)
+ outputs = outputs[:, inputs['input_ids'].shape[1]:]
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+ return response
+
+ def __call__(self, *args, **kwargs):
+ video_path = kwargs.pop('video_path', None)
+ with open(video_path, 'rb') as f:
+ video_data = f.read()
+ response = (self.get_caption(self.prompt, video_data, self.temperature))
+ return response
+
+
+class OpusMtZhEn(BaseCaptionProcessor):
+ def __init__(self, cfg, language='en'):
+ super().__init__(cfg, language=language)
+ self.model_path = cfg.MODEL_PATH
+ self.model_info = {
+ 'device': 'offline',
+ 'model': None,
+ 'tokenizer': None
+ }
+
+ def load_model(self):
+ is_flg, msg = super().load_model()
+ if not is_flg:
+ return is_flg, msg
+ if self.model_info['device'] == 'offline':
+ model = None
+ try:
+ from transformers import MarianMTModel, AutoTokenizer
+ local_model_dir = FS.get_dir_to_local_dir(self.model_path)
+ tokenizer = AutoTokenizer.from_pretrained(
+ local_model_dir
+ )
+ model = MarianMTModel.from_pretrained(
+ local_model_dir
+ ).to(we.device_id)
+ except Exception as e:
+ if model is not None:
+ del model
+ return False, f"Load model error '{e}'"
+ self.model_info['device'] = model.device
+ self.model_info['model'] = model
+ self.model_info['tokenizer'] = tokenizer
+ elif self.model_info['device'] == 'cpu':
+ try:
+ self.model_info['model'].to(we.device_id)
+ self.model_info['device'] = we.device_id
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ except Exception as e:
+ del self.model_info['model']
+ self.model_info['model'] = None
+ self.model_info['device'] = 'offline'
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return False, f"Load model error '{e}'"
+ return True, ''
+
+ def unload_model(self):
+ super().unload_model()
+ if self.delete_instance:
+ self.model_info['device'] = 'offline'
+ if self.model_info['model'] is not None:
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ del self.model_info['model']
+ self.model_info['model'] = None
+ elif (isinstance(self.model_info['device'], numbers.Number)
+ or str(self.model_info['device']).startswith('cuda')):
+ self.model_info['device'] = 'cpu'
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return True, ''
+
+ def get_caption(self, data):
+ model = self.model_info['model']
+ tokenizer = self.model_info['tokenizer']
+ batch = tokenizer(data, return_tensors="pt").to(we.device_id)
+ generated_ids = model.generate(**batch)
+ translated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ return translated_text
+
+ def __call__(self, *args, **kwargs):
+ caption = kwargs.pop('caption', None)
+ response = (self.get_caption(caption))
+ return response
+
+
+class OpusMtEnZh(BaseCaptionProcessor):
+ def __init__(self, cfg, language='en'):
+ super().__init__(cfg, language=language)
+ self.model_path = cfg.MODEL_PATH
+ self.model_info = {
+ 'device': 'offline',
+ 'model': None,
+ 'tokenizer': None
+ }
+
+ def load_model(self):
+ is_flg, msg = super().load_model()
+ if not is_flg:
+ return is_flg, msg
+ if self.model_info['device'] == 'offline':
+ model = None
+ try:
+ from transformers import MarianMTModel, AutoTokenizer
+ local_model_dir = FS.get_dir_to_local_dir(self.model_path)
+ tokenizer = AutoTokenizer.from_pretrained(
+ local_model_dir
+ )
+ model = MarianMTModel.from_pretrained(
+ local_model_dir
+ ).to(we.device_id)
+ except Exception as e:
+ if model is not None:
+ del model
+ return False, f"Load model error '{e}'"
+ self.model_info['device'] = model.device
+ self.model_info['model'] = model
+ self.model_info['tokenizer'] = tokenizer
+ elif self.model_info['device'] == 'cpu':
+ try:
+ self.model_info['model'].to(we.device_id)
+ self.model_info['device'] = we.device_id
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ except Exception as e:
+ del self.model_info['model']
+ self.model_info['model'] = None
+ self.model_info['device'] = 'offline'
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return False, f"Load model error '{e}'"
+ return True, ''
+
+ def unload_model(self):
+ super().unload_model()
+ if self.delete_instance:
+ self.model_info['device'] = 'offline'
+ if self.model_info['model'] is not None:
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ del self.model_info['model']
+ self.model_info['model'] = None
+ elif (isinstance(self.model_info['device'], numbers.Number)
+ or str(self.model_info['device']).startswith('cuda')):
+ self.model_info['device'] = 'cpu'
+ self.model_info['model'] = self.model_info['model'].to('cpu')
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ return True, ''
+
+ def get_caption(self, data):
+ model = self.model_info['model']
+ tokenizer = self.model_info['tokenizer']
+ batch = tokenizer(data, return_tensors="pt").to(we.device_id)
+ generated_ids = model.generate(**batch)
+ translated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ return translated_text
+
+ def __call__(self, *args, **kwargs):
+ caption = kwargs.pop('caption', None)
+ response = (self.get_caption(caption))
+ return response
\ No newline at end of file
diff --git a/scepter/studio/preprocess/utils/txt2vid_data_card.py b/scepter/studio/preprocess/utils/txt2vid_data_card.py
new file mode 100644
index 0000000..9c3af36
--- /dev/null
+++ b/scepter/studio/preprocess/utils/txt2vid_data_card.py
@@ -0,0 +1,347 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import csv
+import os
+import shutil
+import time
+import decord
+from tqdm import tqdm
+import gradio as gr
+
+from scepter.modules.utils.directory import get_md5
+from scepter.modules.utils.file_system import FS
+from scepter.studio.preprocess.caption_editor_ui.component_names import \
+ Text2VideoDataCardName
+from scepter.studio.preprocess.utils.data_card import (BaseDataCard, find_prefix)
+
+
+class Text2VideoDataCard(BaseDataCard):
+ def __init__(self,
+ dataset_folder,
+ dataset_name=None,
+ src_file=None,
+ surfix=None,
+ user_name='admin',
+ language='en'
+ ):
+ super().__init__(dataset_folder,
+ dataset_name=dataset_name,
+ user_name=user_name)
+ self.meta['task_type'] = 'txt2vid'
+ self.components_name = Text2VideoDataCardName(language)
+ if self.new_dataset:
+ if surfix == '.zip':
+ file_list = self.load_from_zip(src_file, dataset_folder,
+ self.local_dataset_folder)
+ elif surfix in ['.txt', '.csv']:
+ file_list = self.load_from_list(src_file, dataset_folder,
+ self.local_dataset_folder)
+ elif surfix is None:
+ file_list = []
+ else:
+ raise gr.Error(
+ f'{self.components_name.illegal_data_err2} {surfix}')
+
+ is_flag = FS.put_dir_from_local_dir(self.local_dataset_folder,
+ dataset_folder,
+ multi_thread=True)
+ if not is_flag:
+ raise gr.Error(f'{self.components_name.illegal_data_err3}')
+
+ self.meta['cursor'] = 0 if len(file_list) > 0 else -1
+ self.meta['file_list'] = file_list
+ self.update_dataset()
+
+ def apply_changes(self):
+ edit_index_list = self.edit_list
+ for index in edit_index_list:
+ one_data = self.data[index]
+ self.data[index]['caption'] = one_data['edit_caption']
+
+ self.update_dataset()
+ return True, ''
+
+ def load_from_list(self, save_file, dataset_folder, local_dataset_folder):
+ file_list = []
+ videos_folder = os.path.join(local_dataset_folder, 'videos')
+ os.makedirs(videos_folder, exist_ok=True)
+ with FS.get_from(save_file) as local_path:
+ with open(local_path, 'r') as f:
+ for line in tqdm(f):
+ line = line.strip()
+ if line == '':
+ continue
+ try:
+ src_video_path, caption = line.split(
+ '#;#', 1)
+ except Exception:
+ try:
+ src_video_path, caption = line.split(
+ ',', 1)
+ except Exception:
+ raise gr.Error(
+ self.components_name.illegal_data_err1)
+ relative_path = os.path.join(
+ 'videos', f'{get_md5(src_video_path)[:18]}_{int(time.time())}.mp4')
+ video_path = os.path.join(dataset_folder, relative_path)
+ FS.get_from(src_video_path, local_path=video_path)
+ is_legal, new_path, prefix = find_prefix(src_video_path)
+ w, h, fps, duration = self.get_video_meta(video_path)
+
+ file_list.append({
+ 'video_path':
+ video_path,
+ 'relative_path':
+ relative_path,
+ 'width':
+ w,
+ 'height':
+ h,
+ 'fps':
+ fps,
+ 'duration':
+ duration,
+ 'caption':
+ caption,
+ 'edit_caption':
+ caption,
+ 'prefix':
+ prefix
+ })
+ return file_list
+
+ def load_from_zip(self, save_file, data_folder, local_dataset_folder):
+ with FS.get_from(save_file) as local_path:
+ res = os.popen(
+ f"unzip -o '{local_path}' -d '{local_dataset_folder}'")
+ res = res.readlines()
+ if not os.path.exists(local_dataset_folder):
+ raise gr.Error(f'Unzip {save_file} failed {str(res)}')
+ file_folder = None
+ train_list = None
+ hit_dir = None
+ raw_list = {}
+ mac_osx = os.path.join(local_dataset_folder, '__MACOSX')
+ if os.path.exists(mac_osx):
+ res = os.popen(f"rm -rf '{mac_osx}'")
+ res = res.readlines()
+ for one_dir in FS.walk_dir(local_dataset_folder, recurse=False):
+ if one_dir.endswith('__MACOSX'):
+ res = os.popen(f"rm -rf '{one_dir}'")
+ res = res.readlines()
+ continue
+ if FS.isdir(one_dir):
+ if one_dir.endswith('videos') or one_dir.endswith('videos/'):
+ file_folder = one_dir
+ hit_dir = one_dir
+ else:
+ sub_dir = FS.walk_dir(one_dir)
+ for one_s_dir in sub_dir:
+ if FS.isdir(one_s_dir) and one_s_dir.split(
+ one_dir)[1].replace('/', '') == 'videos':
+ file_folder = one_s_dir
+ hit_dir = one_dir
+ if FS.isfile(one_s_dir) and one_s_dir.split(
+ one_dir)[1].replace('/', '') == 'train.csv':
+ train_list = one_s_dir
+ if file_folder is not None and train_list is not None:
+ break
+ elif one_dir.endswith('train.csv'):
+ train_list = one_dir
+ else:
+ continue
+ if file_folder is not None and train_list is not None:
+ break
+ if file_folder is None and len(raw_list) < 1:
+ raise gr.Error(
+ "video doesn't exist, or nothing exists in your zip")
+
+ if train_list is None:
+ raise gr.Error("pair list doesn't exist")
+ new_file_folder = f'{local_dataset_folder}/videos'
+ os.makedirs(new_file_folder, exist_ok=True)
+
+ if file_folder is not None:
+ _ = FS.get_dir_to_local_dir(file_folder, new_file_folder)
+
+ if not os.path.exists(new_file_folder):
+ raise gr.Error(f'{str(res)}')
+ new_train_list = f'{local_dataset_folder}/train.csv'
+ res = os.popen(f"mv '{train_list}' '{new_train_list}'")
+ res = res.readlines()
+ if not os.path.exists(new_train_list):
+ raise gr.Error(f'{str(res)}')
+ if not file_folder == hit_dir:
+ try:
+ res = os.popen(f"rm -rf '{hit_dir}/videos/*'")
+ _ = res.readlines()
+ res = os.popen(f"rm -rf '{hit_dir}'")
+ _ = res.readlines()
+
+ except Exception:
+ pass
+ file_list = self.load_train_file(new_train_list)
+ # remove unused data
+ for one_dir in FS.walk_dir(local_dataset_folder):
+ if 'videos' in one_dir or one_dir.endswith(
+ 'file.csv') or one_dir.endswith('train.csv'):
+ continue
+ os.system(f'rm -rf {one_dir}')
+ return file_list
+
+ def load_train_file(self, file_path):
+ base_folder = os.path.dirname(file_path)
+ file_list = []
+ video_set = set()
+
+ with open(file_path, 'r') as f:
+ reader = csv.reader(f)
+ for row in reader:
+ if len(row) == 2:
+ src_video_path, prompt = row[0], row[1]
+ else:
+ return gr.Error(self.components_name.illegal_data_err2)
+ if src_video_path == 'Target:FILE':
+ continue
+
+ local_video_path = os.path.join(base_folder, src_video_path)
+ w, h, fps, duration = self.get_video_meta(local_video_path)
+ if src_video_path in video_set:
+ src_video_path, surfix = os.path.splitext(src_video_path)
+ src_video_path = f'{src_video_path}_{int(time.time() * 100)}{surfix}'
+ new_local_video_path = os.path.join(
+ base_folder, src_video_path)
+ self.copy_video(src_video_path, new_local_video_path)
+ video_set.add(src_video_path)
+
+ file_list.append({
+ 'video_path':
+ local_video_path,
+ 'relative_path':
+ src_video_path,
+ 'width':
+ w,
+ 'height':
+ h,
+ 'fps':
+ fps,
+ "duration":
+ duration,
+ 'caption':
+ prompt,
+ 'edit_caption':
+ prompt,
+ 'prefix':
+ ''
+ })
+ return file_list
+
+ def write_train_file(self):
+ file_list = self.meta['file_list']
+ with open(self.local_train_file, 'w') as f:
+ writer = csv.writer(f)
+ writer.writerow(['Target:FILE', 'Prompt'])
+ for one_file in file_list:
+ relative_file = one_file['relative_path']
+ if relative_file.startswith('/'):
+ relative_file = relative_file[1:]
+ writer.writerow([relative_file, one_file['caption'].
+ strip().replace("\n", "")])
+ FS.put_object_from_local_file(self.local_train_file, self.train_file)
+
+ def write_data_file(self):
+ file_list = self.meta['file_list']
+ with open(self.local_save_file_list, 'w') as f:
+ for one_file in file_list:
+ f.write('{}#;#{}#;#{}#;#{}\n'.format(one_file['relative_path'],
+ one_file['width'],
+ one_file['height'],
+ one_file['caption'].strip().replace("\n", ""))
+ )
+ FS.put_object_from_local_file(self.local_save_file_list,
+ self.save_file_list)
+
+ def add_record(self, video, caption, **kwargs):
+ local_work_dir = self.meta['local_work_dir']
+ work_dir = self.meta['work_dir']
+
+ save_folder = os.path.join(local_work_dir, 'videos')
+ os.makedirs(save_folder, exist_ok=True)
+ w, h, fps, duration = self.get_video_meta(video)
+
+ relative_path = os.path.join(
+ 'videos', f'{get_md5(video)[:18]}_{int(time.time())}.mp4')
+ video_path = os.path.join(work_dir, relative_path)
+ local_video_path = os.path.join(local_work_dir, relative_path)
+ self.copy_video(video, local_video_path)
+
+ self.data.append({
+ 'video_path': video_path,
+ 'relative_path': relative_path,
+ 'width': w,
+ 'height': h,
+ 'fps': fps,
+ 'duration': duration,
+ 'caption': caption,
+ 'edit_caption': caption,
+ 'prefix': ''
+ })
+
+ self.set_cursor(len(self.meta['file_list']) - 1)
+ self.update_dataset()
+ return True
+
+ def copy_video(self, source_path, target_path):
+ if not os.path.isfile(source_path):
+ raise gr.Error('Video path not exist.')
+ try:
+ shutil.copy2(source_path, target_path)
+ except Exception as e:
+ raise gr.Error(str(e))
+
+ def get_video_meta(self, video):
+ video_reader = decord.VideoReader(video)
+ w = video_reader[0].shape[1]
+ h = video_reader[0].shape[0]
+ fps = video_reader.get_avg_fps()
+ video_length = len(video_reader)
+ duration = video_length / fps
+
+ return w, h, fps, duration
+
+ def delete_record(self):
+ if len(self) < 1:
+ raise gr.Error(self.components_name.delete_err1)
+ current_file = self.data.pop(self.cursor)
+ self.set_cursor(self.cursor - 1)
+ local_file = os.path.join(self.meta['local_work_dir'],
+ current_file['relative_path'])
+ try:
+ os.remove(local_file)
+ except Exception:
+ print(f'remove file {local_file} error')
+
+ if self.cursor >= len(self.meta['file_list']):
+ self.set_cursor(0)
+ if self.cursor < 0:
+ self.set_cursor(len(self) - 1)
+ if len(self.meta['file_list']) == 0:
+ self.set_cursor(-1)
+ self.update_dataset()
+
+ def export_zip(self, export_folder):
+ self.update_dataset()
+ zip_path = os.path.join(export_folder, f'{self.dataset_name}.zip')
+ local_zip, _ = FS.map_to_local(zip_path)
+ os.makedirs(os.path.dirname(local_zip), exist_ok=True)
+ res = os.popen(
+ f"cd '{self.local_work_dir}' && mkdir -p '{self.dataset_name}' "
+ f"&& cp -rf videos '{self.dataset_name}/videos' "
+ f"&& cp -rf train.csv '{self.dataset_name}/train.csv' "
+ f"&& zip -r '{os.path.abspath(local_zip)}' '{self.dataset_name}'/* "
+ f"&& rm -rf '{self.dataset_name}'")
+ print(res.readlines())
+ FS.put_object_from_local_file(local_zip, zip_path)
+ if not FS.exists(zip_path):
+ raise gr.Error(self.components_name.export_zip_err1)
+ return local_zip
diff --git a/scepter/studio/self_train/scripts/trainer.py b/scepter/studio/self_train/scripts/trainer.py
index 4ca7924..d7f0a85 100644
--- a/scepter/studio/self_train/scripts/trainer.py
+++ b/scepter/studio/self_train/scripts/trainer.py
@@ -44,18 +44,22 @@ def kill_job(pid):
class Trainer():
- def __init__(self, run_script, status_message):
+ def __init__(self, run_script, status_message, visible_gpus):
self.run_script = run_script
self.status_message = status_message
+ self.visible_gpus = visible_gpus
self.proc = None
def __call__(self, task_name):
torch.cuda.empty_cache()
error_folder = './error_logs'
os.makedirs(error_folder, exist_ok=True)
+
self.status_message.error_log = f'{error_folder}/{int(time.time())}.log'
cmd = f'PYTHONPATH=. python {self.run_script} ' \
f'--cfg={task_name}/train.yaml 2> {self.status_message.error_log}'
+ if self.visible_gpus is not None and len(self.visible_gpus) > 0:
+ cmd = f'CUDA_VISIBLE_DEVICES={",".join([str(i) for i in self.visible_gpus])} ' + cmd
# cmd = [f"python {self.run_script}"]
print(cmd)
try:
@@ -87,6 +91,7 @@ class TrainManager():
self.runing_tasks = {}
self.run_script = run_script
self.work_dir = work_dir
+ self.visible_gpus = list(range(torch.cuda.device_count()))
def task_dispatch():
while True:
@@ -143,7 +148,7 @@ class TrainManager():
task_name = self.task_queue.pop(0)
print(f'start task {task_name}')
status_message = TaskStatus()
- train_ins = Trainer(self.run_script, status_message)
+ train_ins = Trainer(self.run_script, status_message, self.visible_gpus)
train_thread = threading.Thread(target=train_ins,
args=(os.path.join(
self.work_dir,
@@ -174,11 +179,18 @@ class TrainManager():
self.task_manage = threading.Thread(target=task_dispatch, daemon=True)
self.task_manage.start()
+ def set_gpus(self, gpus=None):
+ if gpus is None:
+ self.visible_gpus = list(range(torch.cuda.device_count()))
+ else:
+ self.visible_gpus = list(gpus)
+
def check_memory(self):
# Check Cuda Memory
+ visible_gpus = self.visible_gpus
mem_msg = ''
if torch.cuda.is_available():
- for device_id in range(torch.cuda.device_count()):
+ for device_id in visible_gpus:
free_mem, total_mem = torch.cuda.mem_get_info(device_id)
free_mem = free_mem / (1024**3)
total_mem = total_mem / (1024**3)
diff --git a/scepter/studio/self_train/self_train_ui/component_names.py b/scepter/studio/self_train/self_train_ui/component_names.py
index d809d4b..84757f4 100644
--- a/scepter/studio/self_train/self_train_ui/component_names.py
+++ b/scepter/studio/self_train/self_train_ui/component_names.py
@@ -74,12 +74,13 @@ class TrainerUIName():
self.task_choices = ['Text2Image', 'Image Editing']
self.data_task_map = {
'scepter_txt2img': None,
- 'scepter_img2img': 'edit'
+ 'scepter_img2img': 'edit',
+ 'scepter_txt2vid': 'dit'
}
if language == 'en':
self.user_direction = '''
### User Guide
- - Data: Data preparation is done through the Data Manager. (zip example: [3D](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets/3D_example_csv.zip))
+ - Data: Data preparation is done through the Data Manager. (zip example: [3D](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets/3D_example_csv.zip), txt example(Support only video data): [txt](https://modelscope.cn/models/iic/scepter/resolve/master/datasets/video_example.txt))
- Parameters: You can try modifying the related parameters.
- Training: Click on [Start Training].
- Testing: After completing the training, click [Go to inference].
@@ -93,14 +94,20 @@ class TrainerUIName():
self.data_source_choices = [
'Dataset zip', 'MaaS Dataset', 'Dataset Management'
]
+ self.illegal_data_err = (
+ 'The list supports only "," or "#;#" as delimiters. '
+ 'The two columns represent video path and description, '
+ 'respectively.')
self.data_source_value = 'Dataset zip'
self.data_source_name = 'Data Source'
self.data_type_map = {
'scepter_txt2img': 'Text2Image Generation',
- 'scepter_img2img': 'Image Edit Generation'
+ 'scepter_img2img': 'Image Edit Generation',
+ 'scepter_txt2vid': 'Text2Video Generation'
}
self.data_type_choices = list(self.data_type_map.keys())
self.data_type_value = 'scepter_txt2img'
+ self.data_type_value_video = 'scepter_txt2vid'
self.data_type_name = 'Data Type'
self.ori_data_name = 'Data Name'
# Supports MaaS dataset/local/HTTP Zip package
@@ -120,8 +127,8 @@ class TrainerUIName():
self.base_model = 'Base Model'
self.tuner_name = 'Tuner Method'
self.base_model_revision = 'Model Version Number'
- self.resolution_height = 'Train Image Height'
- self.resolution_width = 'Train Image Width'
+ self.resolution_height = 'Train Image or Video Height'
+ self.resolution_width = 'Train Image or Video Width'
self.resolution_height_max = 'Resolution Height Max'
self.resolution_width_max = 'Resolution Width Max'
self.train_epoch = 'Total Training Epochs'
@@ -142,6 +149,8 @@ class TrainerUIName():
self.bucket_resolution_steps = 'Bucket Resolution Steps'
self.bucket_no_upscale = 'Bucket No Upscale'
self.bucket_no_upscale_ins = 'Disable Automatic Image Upscaling'
+ self.accumulate_step = 'Accumulate Step'
+ self.gpus = 'Select GPUs'
# Error or Warning
self.training_err1 = 'CUDA is unavailable.'
self.training_err2 = 'Currently insufficient VRAM, training failed!'
@@ -153,7 +162,7 @@ class TrainerUIName():
elif language == 'zh':
self.user_direction = '''
### 使用说明
- - 数据: 通过数据管理器进行数据的准备（ZIP样例：[3D](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets/3D_example_csv.zip)）
+ - 数据: 通过数据管理器进行数据的准备（ZIP样例：[3D](https://www.modelscope.cn/api/v1/models/iic/scepter/repo?Revision=master&FilePath=datasets/3D_example_csv.zip)，txt样例（仅支持视频数据）：[txt](https://modelscope.cn/models/iic/scepter/resolve/master/datasets/video_example.txt)）
- 参数: 可尝试进行相关参数的修改
- 训练: 点击【开始训练】
- 测试: 完成训练后点击【使用模型】
@@ -161,14 +170,17 @@ class TrainerUIName():
- 对于大规模数据的处理和训练，建议使用命令行形式
''' # noqa
self.data_source_choices = ['数据集zip', 'MaaS数据集', '数据管理器']
+ self.illegal_data_err = '列表只支持,或#;#作为分割符，两列分别为视频路径/描述'
self.data_source_value = '数据集zip'
self.data_source_name = '数据集来源'
self.data_type_map = {
'scepter_txt2img': '文生图数据',
- 'scepter_img2img': '图像编辑（图生图）数据'
+ 'scepter_img2img': '图像编辑（图生图）数据',
+ 'scepter_txt2vid': '文生视频数据'
}
self.data_type_choices = list(self.data_type_map.keys())
self.data_type_value = 'scepter_txt2img'
+ self.data_type_value_video = 'scepter_txt2vid'
self.data_type_name = '数据类型'
self.ori_data_name = '数据集名称'
self.ms_data_name_place_hold = '请使用数据管理器导入' # '支持MaaS数据集/本地/Http Zip包'
@@ -187,8 +199,8 @@ class TrainerUIName():
self.base_model = '基础模型'
self.tuner_name = '微调方法'
self.base_model_revision = '模型版本号'
- self.resolution_height = '训练图片高度'
- self.resolution_width = '训练图片宽度'
+ self.resolution_height = '训练图片或视频高度'
+ self.resolution_width = '训练图片或视频宽度'
self.resolution_height_max = '最大训练高度'
self.resolution_width_max = '最大训练宽度'
self.train_epoch = '总训练轮数'
@@ -208,6 +220,8 @@ class TrainerUIName():
self.bucket_resolution_steps = '分桶分辨率步长'
self.bucket_no_upscale = '分桶分辨率不做放大'
self.bucket_no_upscale_ins = '禁止图片分辨率上采样'
+ self.accumulate_step = '梯度累积数量'
+ self.gpus = '选择GPU'
# Error or Warning
self.training_err1 = 'CUDA不可用.'
self.training_err2 = '目前显存不足，训练失败！'
diff --git a/scepter/studio/self_train/self_train_ui/model_ui.py b/scepter/studio/self_train/self_train_ui/model_ui.py
index 7cac932..1b87e08 100644
--- a/scepter/studio/self_train/self_train_ui/model_ui.py
+++ b/scepter/studio/self_train/self_train_ui/model_ui.py
@@ -59,13 +59,16 @@ class ModelUI(UIBase):
status_file = os.path.join(self.work_dir, one_dir,
'status.json')
if FS.exists(status_file):
- status = json.load(open(status_file, 'r'))
+ try:
+ status = json.load(open(status_file, 'r'))
+ except:
+ continue
status['model_name'] = one_dir
have_model_list.append(status)
have_model_list.sort(key=lambda x: x['start_time'])
self.user_level_model_list[user_name] = [
v['model_name'] for v in have_model_list
- ][:100]
+ ]
def get_ckpt_list(self, output_model):
all_ckpt_list = []
@@ -177,11 +180,12 @@ class ModelUI(UIBase):
def model_name_change(model_name):
if model_name is None:
- return '', gr.Column(), '', []
+ return '', gr.Column(), None, []
message = trainer_ui.trainer_ins.get_log(model_name)
status = trainer_ui.trainer_ins.get_status(model_name)
ckpt_list = self.get_ckpt_list(model_name)
- ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else ''
+ ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else None
+ ckpt_list = ckpt_list if isinstance(ckpt_list, list) and len(ckpt_list) > 0 else None
if ckpt_value is not None and len(ckpt_value) > 0:
gallery_value = self.get_gallery_list(model_name, ckpt_value)
else:
@@ -264,13 +268,14 @@ class ModelUI(UIBase):
message = trainer_ui.trainer_ins.get_log(model_name)
status = trainer_ui.trainer_ins.get_status(model_name)
ckpt_list = self.get_ckpt_list(model_name)
- ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else ''
+ ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else None
+ ckpt_list = ckpt_list if isinstance(ckpt_list, list) and len(ckpt_list) > 0 else None
ret_gallery = ckpt_name_change(model_name, ckpt_value)
+ model_list = self.user_level_model_list.get(login_user_name, [])
self.load_history(login_user_name)
return (message, gr.Column(visible=status in ('running',
'success')),
- gr.Dropdown(choices=self.user_level_model_list.get(
- login_user_name, []),
+ gr.Dropdown(choices=model_list,
value=model_name),
gr.Dropdown(choices=ckpt_list,
value=ckpt_value), ret_gallery)
@@ -330,7 +335,8 @@ class ModelUI(UIBase):
message = trainer_ui.trainer_ins.get_log(model_name)
status = trainer_ui.trainer_ins.get_status(model_name)
ckpt_list = self.get_ckpt_list(model_name)
- ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else ''
+ ckpt_value = ckpt_list[-1] if len(ckpt_list) > 0 else None
+ ckpt_list = ckpt_list if isinstance(ckpt_list, list) and len(ckpt_list) > 0 else None
ret_gallery = ckpt_name_change(model_name, ckpt_value)
self.load_history(login_user_name)
return (message, gr.Column(visible=status in ('running',
@@ -553,7 +559,7 @@ class ModelUI(UIBase):
if len(model_list) > 0:
model_name = model_list[-1]
else:
- model_name = ''
+ model_name = None
return gr.Dropdown(choices=model_list, value=model_name)
manager.user_name.change(model_name_change,
diff --git a/scepter/studio/self_train/self_train_ui/trainer_ui.py b/scepter/studio/self_train/self_train_ui/trainer_ui.py
index 948d5da..c84d3d0 100644
--- a/scepter/studio/self_train/self_train_ui/trainer_ui.py
+++ b/scepter/studio/self_train/self_train_ui/trainer_ui.py
@@ -5,12 +5,17 @@ import datetime
import json
import os
import random
+import time
from collections import OrderedDict
+from tqdm import tqdm
+
+import decord
import gradio as gr
import scepter
import torch
import yaml
+from scepter.modules.utils.directory import get_md5
from scepter.modules.utils.file_system import FS
from scepter.studio.self_train.scripts.trainer import TrainManager
from scepter.studio.self_train.self_train_ui.component_names import \
@@ -84,6 +89,7 @@ class TrainerUI(UIBase):
self.current_train_model = None
self.trainer_ins = TrainManager(self.run_script, self.work_dir_pre)
self.component_names = TrainerUIName(language=language)
+ self.save_file_local_path = cfg.SAVE_FILE_LOCAL_PATH
self.h_level_dict = {}
for hw_tuple in self.train_para_data.RESOLUTIONS.get('VALUES', []):
@@ -107,7 +113,7 @@ class TrainerUI(UIBase):
choices=self.component_names.data_source_choices,
value=self.component_names.data_source_value,
label=self.component_names.data_source_name,
- interactive=False)
+ interactive=True)
self.data_type = gr.Dropdown(
choices=[
self.component_names.data_type_map[key] for key
@@ -116,20 +122,20 @@ class TrainerUI(UIBase):
value=self.component_names.data_type_map[
self.component_names.data_type_value],
label=self.component_names.data_type_name,
- interactive=False)
+ interactive=True)
self.ori_data_name = gr.Textbox(
label=self.component_names.ori_data_name,
max_lines=1,
placeholder=self.component_names.ori_data_name,
- interactive=False)
+ interactive=True)
self.ms_data_name = gr.Textbox(
label=' or '.join(
self.component_names.data_source_choices),
max_lines=1,
placeholder=self.component_names.
ms_data_name_place_hold,
- visible=False,
- interactive=False)
+ visible=True,
+ interactive=True)
with gr.Group(visible=False) as self.ms_data_box:
with gr.Row():
self.ms_data_space = gr.Textbox(
@@ -199,8 +205,7 @@ class TrainerUI(UIBase):
visible=lora_visible) as self.lora_param:
self.lora_alpha = gr.Number(
label='LoRA Alpha',
- value=self.para_data.get(
- 'lora_alpha', 4),
+ value=self.para_data.get('lora_alpha', 4),
interactive=True)
self.lora_rank = gr.Number(
label='LoRA Rank',
@@ -310,6 +315,19 @@ class TrainerUI(UIBase):
precision=0,
interactive=True)
+ with gr.Row():
+ self.accumulate_step = gr.Number(
+ label=self.component_names.accumulate_step,
+ value=self.para_data.get('ACCUMULATE_STEP', 1),
+ precision=0,
+ interactive=True)
+ self.gpus = gr.Dropdown(
+ choices=list(range(torch.cuda.device_count())),
+ value=list(range(torch.cuda.device_count())),
+ label=self.component_names.gpus,
+ multiselect=True,
+ interactive=True)
+
with gr.Row():
self.prompt_prefix = gr.Text(
label=self.component_names.prompt_prefix,
@@ -355,6 +373,15 @@ class TrainerUI(UIBase):
self.component_names.data_type_value], 'damo',
'style_custom_dataset', 'style_custom_dataset',
'3D'
+ ],
+ [
+ self.component_names.data_source_choices[2],
+ self.component_names.data_type_map[
+ self.component_names.data_type_value_video],
+ '',
+ 'https://modelscope.cn/models/iic/scepter/resolve/master/datasets/video_example.txt', # noqa
+ 'video_example_txt',
+ ''
]
],
inputs=[
@@ -439,6 +466,7 @@ class TrainerUI(UIBase):
eval_prompts = [] if is_edit else self.train_para_data.get(
'EVAL_PROMPTS', [])
eval_prompts = ret_data.get('EVAL_PROMPTS', eval_prompts)
+
return ret_data.get('EPOCHS', 10), \
ret_data.get('LEARNING_RATE', 0.0001), \
ret_data.get('SAVE_INTERVAL', 10), \
@@ -613,7 +641,8 @@ class TrainerUI(UIBase):
lora_alpha, lora_rank, text_lora_alpha, text_lora_rank,
sce_ratio, enable_resolution_bucket,
min_bucket_resolution, max_bucket_resolution,
- bucket_resolution_steps, bucket_no_upscale, user_name):
+ bucket_resolution_steps, bucket_no_upscale,
+ accumulate_step, gpus, user_name):
# Check Cuda
if not torch.cuda.is_available() and not self.is_debug:
raise gr.Error(self.component_names.training_err1)
@@ -621,7 +650,7 @@ class TrainerUI(UIBase):
if work_name == 'custom' or work_name is None or work_name == '':
raise gr.Error(self.component_names.training_err4)
work_dir = os.path.join(self.work_dir_pre, work_name)
- login_user_name = user_name
+
self.current_train_model = work_name
if os.path.exists(work_dir) or os.path.exists(
f'.flag/{work_name}.tmp'):
@@ -655,7 +684,7 @@ class TrainerUI(UIBase):
if ms_data_name is None:
raise gr.Error(self.component_names.training_err3)
- def prepare_train_data(data_cfg):
+ def prepare_train_image_data(data_cfg):
data_cfg['BATCH_SIZE'] = int(train_batch_size)
data_cfg['PROMPT_PREFIX'] = prompt_prefix
data_cfg['REPLACE_KEYWORDS'] = replace_keywords
@@ -737,8 +766,12 @@ class TrainerUI(UIBase):
if os.path.exists(local_data_dir) and os.path.exists(
local_file_list):
data_cfg.update({
- 'NAME': 'ImageTextPairDataset' if data_cfg['NAME'] == 'ImageTextPairMSDataset' else data_cfg['NAME'],
- 'ENABLE_RESOLUTION_BUCKET': enable_resolution_bucket,
+ 'NAME':
+ 'ImageTextPairDataset'
+ if data_cfg['NAME'] == 'ImageTextPairMSDataset'
+ else data_cfg['NAME'],
+ 'ENABLE_RESOLUTION_BUCKET':
+ enable_resolution_bucket,
'SAMPLER': {
'NAME':
'ResolutionBatchSampler',
@@ -761,7 +794,8 @@ class TrainerUI(UIBase):
'BUCKET_NO_UPSCALE':
bucket_no_upscale
},
- 'DATA_NUM': data_num
+ 'DATA_NUM':
+ data_num
})
if 'TRANSFORMS' in data_cfg:
for trans in data_cfg['TRANSFORMS']:
@@ -776,6 +810,69 @@ class TrainerUI(UIBase):
return data_cfg
+ def prepare_train_video_data(data_cfg):
+ if ms_data_name.startswith('http') and (
+ '.txt' in ms_data_name or '.csv' in ms_data_name):
+ data_name = get_data_from_list()
+ else:
+ data_name = os.path.join(ms_data_name, 'file.txt')
+
+ data_cfg['BATCH_SIZE'] = int(train_batch_size)
+ data_cfg['PROMPT_PREFIX'] = prompt_prefix
+ if data_cfg['NAME'] in ['VideoGenDataset']:
+ data_cfg['SAMPLER']['SUB_SAMPLERS'][0][
+ 'PATH_PREFIX'] = os.path.dirname(data_name)
+ data_cfg['SAMPLER']['SUB_SAMPLERS'][0][
+ 'INDEX_FILE'] = data_name
+ elif data_cfg['NAME'] in ['VideoGenDatasetOTF']:
+ data_cfg['PATH_PREFIX'] = os.path.dirname(data_name)
+ data_cfg['DATA_FILE'] = data_name
+ else:
+ raise Exception('Unsupported data type {}'.format(
+ data_cfg['NAME']))
+ return data_cfg
+
+ def get_data_from_list():
+ file_list = []
+ file = FS.get_from(ms_data_name)
+ with FS.get_from(file) as local_path:
+ with open(local_path, 'r') as f:
+ for line in tqdm(f):
+ line = line.strip()
+ if line == '':
+ continue
+ try:
+ src_video_path, caption = line.split('#;#', 1)
+ except Exception:
+ try:
+ src_video_path, caption = line.split(
+ ',', 1)
+ except Exception:
+ raise gr.Error(
+ self.component_names.illegal_data_err)
+ relative_path = os.path.join(
+ 'videos',
+ f'{get_md5(src_video_path)[:18]}_{int(time.time())}.mp4'
+ )
+ video_path = os.path.join(
+ self.save_file_local_path, ori_data_name,
+ relative_path)
+ local_path = FS.get_from(src_video_path,
+ local_path=video_path)
+ video_reader = decord.VideoReader(local_path)
+ w = video_reader[0].shape[1]
+ h = video_reader[0].shape[0]
+ file_list.append('{}#;#{}#;#{}#;#{}\n'.format(
+ relative_path, w, h, caption))
+ local_save_file_list = os.path.join(self.save_file_local_path,
+ ori_data_name, 'file.txt')
+ directory = os.path.dirname(local_save_file_list)
+ os.makedirs(directory, exist_ok=True)
+ FS.delete_object(file)
+ with open(local_save_file_list, 'w') as f:
+ f.writelines(file_list)
+ return local_save_file_list
+
def prepare_eval_data(data_cfg):
data_cfg['PROMPT_PREFIX'] = prompt_prefix
data_cfg['IMAGE_SIZE'] = [
@@ -850,8 +947,14 @@ class TrainerUI(UIBase):
]
cfg['SOLVER']['TUNER'] = tuner_cfg_list
- cfg['SOLVER']['TRAIN_DATA'] = prepare_train_data(
- cfg['SOLVER']['TRAIN_DATA'])
+ if cfg['SOLVER']['TRAIN_DATA']['NAME'] in [
+ 'VideoGenDataset', 'VideoGenDatasetOTF'
+ ]:
+ cfg['SOLVER']['TRAIN_DATA'] = prepare_train_video_data(
+ cfg['SOLVER']['TRAIN_DATA'])
+ else:
+ cfg['SOLVER']['TRAIN_DATA'] = prepare_train_image_data(
+ cfg['SOLVER']['TRAIN_DATA'])
if eval_prompts is not None and len(eval_prompts) > 0:
cfg['SOLVER']['EVAL_DATA'] = prepare_eval_data(
cfg['SOLVER']['EVAL_DATA'])
@@ -868,6 +971,8 @@ class TrainerUI(UIBase):
hook['INTERVAL'] = save_interval
hook['PUSH_TO_HUB'] = push_to_hub
hook['HUB_MODEL_ID'] = hub_model_id
+ if hook['NAME'] == 'BackwardHook':
+ hook['ACCUMULATE_STEP'] = int(accumulate_step)
if 'EVAL_HOOKS' in cfg['SOLVER']:
for hook in cfg['SOLVER']['EVAL_HOOKS']:
if hook['NAME'] == 'ProbeDataHook':
@@ -881,6 +986,7 @@ class TrainerUI(UIBase):
default_flow_style=False)
return cfg_file
+ self.trainer_ins.set_gpus(gpus)
before_kill_inference = self.trainer_ins.check_memory()
if hasattr(manager, 'inference'):
for k, v in manager.inference.pipe_manager.pipeline_level_modules.items(
@@ -892,7 +998,6 @@ class TrainerUI(UIBase):
'dynamic_unload')):
manager.preprocess.dataset_gallery.processors_manager.dynamic_unload(
)
-
after_kill_inference = self.trainer_ins.check_memory()
message = f'GPU info: {before_kill_inference}. \n\n'
message += f'After unloading inference models, the GPU info: {after_kill_inference}. \n\n'
@@ -923,7 +1028,8 @@ class TrainerUI(UIBase):
self.text_lora_alpha, self.text_lora_rank, self.sce_ratio,
self.enable_resolution_bucket, self.min_bucket_resolution,
self.max_bucket_resolution, self.bucket_resolution_steps,
- self.bucket_no_upscale, manager.user_name
+ self.bucket_no_upscale, self.accumulate_step, self.gpus,
+ manager.user_name
],
outputs=[inference_ui.output_model_name],
queue=True)
diff --git a/scepter/tools/webui.py b/scepter/tools/webui.py
index be33021..7e2ba60 100644
--- a/scepter/tools/webui.py
+++ b/scepter/tools/webui.py
@@ -65,6 +65,13 @@ if __name__ == '__main__':
choices=['en', 'zh'],
default='en',
help='Now we only support english(en) and chinese(zh)')
+ parser.add_argument('--tab',
+ dest='tab',
+ choices=['all', 'chatbot'],
+ default='all',
+ help='The tabs will be launched, '
+ 'set [all] to use all tools and set [chatbot] to use chatbot only.')
+
args = parser.parse_args()
if not os.path.exists(args.config):
print(
@@ -88,7 +95,7 @@ if __name__ == '__main__':
if not FS.exists(info['CONFIG']):
raise f"{info['CONFIG']} doesn't exist."
interface = None
- if ifid == 'home':
+ if ifid == 'home' and args.tab in ["all", ifid]:
from scepter.studio.home.home import HomeUI
interface = HomeUI(info['CONFIG'],
@@ -96,7 +103,7 @@ if __name__ == '__main__':
language=args.language,
root_work_dir=config.WORK_DIR)
print('init home page success!')
- if ifid == 'preprocess':
+ if ifid == 'preprocess' and args.tab in ["all", ifid]:
from scepter.studio.preprocess.preprocess import PreprocessUI
interface = PreprocessUI(info['CONFIG'],
@@ -104,7 +111,7 @@ if __name__ == '__main__':
language=args.language,
root_work_dir=config.WORK_DIR)
print('init preprocess success!')
- if ifid == 'self_train':
+ if ifid == 'self_train' and args.tab in ["all", ifid]:
from scepter.studio.self_train.self_train import SelfTrainUI
interface = SelfTrainUI(info['CONFIG'],
@@ -112,14 +119,14 @@ if __name__ == '__main__':
language=args.language,
root_work_dir=config.WORK_DIR)
print('init self-train success!')
- if ifid == 'tuner_manager':
+ if ifid == 'tuner_manager' and args.tab in ["all", ifid]:
from scepter.studio.tuner_manager.tuner_manager import TunerManagerUI
interface = TunerManagerUI(info['CONFIG'],
is_debug=args.debug,
language=args.language,
root_work_dir=config.WORK_DIR)
print('init tuner-manager success!')
- if ifid == 'inference':
+ if ifid == 'inference' and args.tab in ["all", ifid]:
from scepter.studio.inference.inference import InferenceUI
interface = InferenceUI(info['CONFIG'],
@@ -127,7 +134,7 @@ if __name__ == '__main__':
language=args.language,
root_work_dir=config.WORK_DIR)
print('init inference success!')
- if ifid == 'ChatBot':
+ if ifid == 'chatbot' and args.tab in ["all", ifid]:
from scepter.studio.chatbot.chatbot import ChatBotUI
interface = ChatBotUI(info['CONFIG'],
@@ -177,10 +184,13 @@ if __name__ == '__main__':
if len(auth_info) > 0:
demo.load(init_value, outputs=[tab_manager.user_name])
+ allowed_paths = [config['WORK_DIR']]
+ allowed_paths.extend(list(set([fs_cfg['TEMP_DIR'] for fs_cfg in config['FILE_SYSTEM']])) if 'FILE_SYSTEM' in config else [])
demo.queue(status_update_rate=1).launch(
server_name=args.host if args.host else config['HOST'],
server_port=int(args.port) if args.port else config['PORT'],
root_path=config['ROOT'],
show_error=True,
debug=True,
- auth=check_auth if len(auth_info) > 0 else None)
+ auth=check_auth if len(auth_info) > 0 else None,
+ allowed_paths=allowed_paths)
diff --git a/scepter/version.py b/scepter/version.py
index 79b1b0a..8c43cf5 100644
--- a/scepter/version.py
+++ b/scepter/version.py
@@ -1,7 +1,7 @@
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
-__version__ = '1.2.0'
+__version__ = '1.3.0'
version_info = tuple(int(x) for x in __version__.split('.')[0:3])
diff --git a/scepter/workflow/config/ace_0.6b_1024_pro.yaml b/scepter/workflow/config/ace_0.6b_1024_pro.yaml
new file mode 100644
index 0000000..2d99b72
--- /dev/null
+++ b/scepter/workflow/config/ace_0.6b_1024_pro.yaml
@@ -0,0 +1,296 @@
+NAME: ACE_0.6B_1024
+IS_DEFAULT: False
+DEFAULT_PARAS:
+ PARAS:
+ #
+ INPUT:
+ INPUT_IMAGE:
+ INPUT_MASK:
+ TASK:
+ PROMPT: ""
+ NEGATIVE_PROMPT: ""
+ OUTPUT_HEIGHT: 1024
+ OUTPUT_WIDTH: 1024
+ SAMPLER: ddim
+ SAMPLE_STEPS: 20
+ GUIDE_SCALE: 4.5
+ GUIDE_RESCALE: 0.5
+ SEED: -1
+ TAR_INDEX: 0
+ OUTPUT:
+ LATENT:
+ IMAGES:
+ SEED:
+ MODULES_PARAS:
+ FIRST_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: float16
+ INPUT: ["IMAGE"]
+ - NAME: decode
+ DTYPE: float16
+ INPUT: ["LATENT"]
+ #
+ DIFFUSION_MODEL:
+ FUNCTION:
+ - NAME: forward
+ DTYPE: float16
+ INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
+ #
+ COND_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode_list_of_list
+ DTYPE: bfloat16
+ INPUT: ["PROMPT"]
+#
+TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+USE_TEXT_POS_EMBEDDINGS: True
+#
+MODEL:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: ms://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
+#
+MODEL_LOCAL:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: models/scepter/ACE-0.6B-1024px/models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
+#
+MODEL_HF:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: hf://scepter-studio/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
\ No newline at end of file
diff --git a/scepter/workflow/config/ace_0.6b_1024_refiner_pro.yaml b/scepter/workflow/config/ace_0.6b_1024_refiner_pro.yaml
new file mode 100644
index 0000000..d40b62d
--- /dev/null
+++ b/scepter/workflow/config/ace_0.6b_1024_refiner_pro.yaml
@@ -0,0 +1,729 @@
+NAME: ACE_0.6B_1024_REFINER
+IS_DEFAULT: False
+DEFAULT_PARAS:
+ PARAS:
+ #
+ INPUT:
+ INPUT_IMAGE:
+ INPUT_MASK:
+ TASK:
+ PROMPT: ""
+ NEGATIVE_PROMPT: ""
+ OUTPUT_HEIGHT: 1024
+ OUTPUT_WIDTH: 1024
+ SAMPLER: ddim
+ SAMPLE_STEPS: 50
+ GUIDE_SCALE: 4.5
+ GUIDE_RESCALE: 0.5
+ SEED: -1
+ TAR_INDEX: 0
+ OUTPUT:
+ LATENT:
+ IMAGES:
+ SEED:
+ MODULES_PARAS:
+ FIRST_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: float16
+ INPUT: ["IMAGE"]
+ - NAME: decode
+ DTYPE: float16
+ INPUT: ["LATENT"]
+ #
+ DIFFUSION_MODEL:
+ FUNCTION:
+ - NAME: forward
+ DTYPE: float16
+ INPUT: ["SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE"]
+ #
+ COND_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode_list_of_list
+ DTYPE: bfloat16
+ INPUT: ["PROMPT"]
+#
+MODEL:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: ms://iic/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: ms://iic/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
+#
+MODEL_LOCAL:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: models/scepter/ACE-0.6B-1024px/models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: models/scepter/ACE-0.6B-1024px/models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
+#
+MODEL_HF:
+ NAME: LatentDiffusionACE
+ PRETRAINED_MODEL:
+ IGNORE_KEYS: [ ]
+ SCALE_FACTOR: 0.18215
+ SIZE_FACTOR: 8
+ DECODER_BIAS: 0.5
+ DEFAULT_N_PROMPT: ""
+ TEXT_IDENTIFIER: [ '{image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ USE_TEXT_POS_EMBEDDINGS: True
+ #
+ DIFFUSION:
+ NAME: BaseDiffusion
+ PREDICTION_TYPE: eps
+ MIN_SNR_GAMMA:
+ NOISE_SCHEDULER:
+ NAME: LinearScheduler
+ NUM_TIMESTEPS: 1000
+ BETA_MIN: 0.0001
+ BETA_MAX: 0.02
+ #
+ DIFFUSION_MODEL:
+ NAME: ACE
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/dit/ace_0.6b_1024px.pth
+ IGNORE_KEYS: [ ]
+ PATCH_SIZE: 2
+ IN_CHANNELS: 4
+ HIDDEN_SIZE: 1152
+ DEPTH: 28
+ NUM_HEADS: 16
+ MLP_RATIO: 4.0
+ PRED_SIGMA: True
+ DROP_PATH: 0.0
+ WINDOW_DIZE: 0
+ Y_CHANNELS: 4096
+ MAX_SEQ_LEN: 1024
+ QK_NORM: True
+ USE_GRAD_CHECKPOINT: True
+ ATTENTION_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKL
+ EMBED_DIM: 4
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/vae/vae.bin
+ IGNORE_KEYS: []
+ #
+ ENCODER:
+ NAME: Encoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 4
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5EmbedderHF
+ PRETRAINED_MODEL: hf://scepter-studio/ACE-0.6B-1024px@models/text_encoder/t5-v1_1-xxl/
+ TOKENIZER_PATH: hf://scepter-studio/ACE-0.6B-1024px@models/tokenizer/t5-v1_1-xxl
+ LENGTH: 120
+ T5_DTYPE: bfloat16
+ ADDED_IDENTIFIER: [ '{image}', '{caption}', '{mask}', '{ref_image}', '{image1}', '{image2}', '{image3}', '{image4}', '{image5}', '{image6}', '{image7}', '{image8}', '{image9}' ]
+ CLEAN: whitespace
+ USE_GRAD: False
+
+REFINER_SCALE: 0.4
+#REFINER_PROMPT: "High Resolution, Sharpness, Clarity, Detail Enhancement, Noise Reduction, HD, 4k, Image Restoration, HDR"
+REFINER_PROMPT: ""
+ACE_PROMPT: [
+ "A cute cartoon rabbit holding a whiteboard that says 'ACE Refiner', standing in a sunny meadow filled with flowers, with a big smile and bright colors.",
+ "A beautiful young woman with long flowing hair, wearing a summer dress, holding a whiteboard that reads 'ACE Refiner' while sitting on a park bench surrounded by cherry blossoms.",
+ "An adorable cartoon cat wearing oversized glasses, holding a whiteboard that says 'ACE Refiner', perched on a stack of colorful books in a cozy library setting.",
+ "A charming girl with pigtails, wearing a cute school uniform, enthusiastically holding a whiteboard that has 'ACE Refiner' written on it, in a bright and cheerful classroom full of educational posters.",
+ "A friendly cartoon dog with floppy ears, sitting in front of a doghouse, proudly holding a whiteboard that says 'ACE Refiner', with a playful expression and a blue sky in the background.",
+ "A cute anime girl with big expressive eyes, dressed in a colorful outfit, holding a whiteboard that reads 'ACE Refiner' in a fantastical landscape filled with mythical creatures.",
+ "A vibrant cartoon fox holding a whiteboard that says 'ACE Refiner', standing on a rock by a sparkling stream, surrounded by lush greenery and butterflies.",
+ "A stylish young woman in a business outfit, smiling as she holds a whiteboard written with 'ACE Refiner', in a modern office filled with plants and natural light.",
+ "A cute cartoon unicorn holding a sparkling whiteboard that says 'ACE Refiner', frolicking in a magical forest, with rainbows and stars in the background.",
+ "A happy family, consisting of a cute little girl and her playful puppy, holding a whiteboard that says 'ACE Refiner', together in their backyard on a sunny day."
+]
+
+REFINER_MODEL:
+ NAME: ""
+ IS_DEFAULT: False
+ DEFAULT_PARAS:
+ PARAS:
+ RESOLUTIONS: [ [ 1024, 1024 ] ]
+ INPUT:
+ INPUT_IMAGE:
+ INPUT_MASK:
+ TASK:
+ PROMPT: ""
+ NEGATIVE_PROMPT: ""
+ OUTPUT_HEIGHT: 1024
+ OUTPUT_WIDTH: 1024
+ SAMPLER: flow_euler
+ SAMPLE_STEPS: 50
+ GUIDE_SCALE: 3.5
+ GUIDE_RESCALE: 0.5
+ OUTPUT:
+ LATENT:
+ IMAGES:
+ SEED:
+ MODULES_PARAS:
+ FIRST_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "IMAGE" ]
+ - NAME: decode
+ DTYPE: bfloat16
+ INPUT: [ "LATENT" ]
+ PARAS:
+ SCALE_FACTOR: 1.5305
+ SHIFT_FACTOR: 0.0609
+ SIZE_FACTOR: 8
+ DIFFUSION_MODEL:
+ FUNCTION:
+ - NAME: forward
+ DTYPE: bfloat16
+ INPUT: [ "SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE" ]
+ COND_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "PROMPT" ]
+
+ MODEL:
+ DIFFUSION:
+ NAME: DiffusionFluxRF
+ PREDICTION_TYPE: raw
+ NOISE_SCHEDULER:
+ NAME: FlowMatchSigmaScheduler
+ WEIGHTING_SCHEME: logit_normal
+ SHIFT: 3.0
+ LOGIT_MEAN: 0.0
+ LOGIT_STD: 1.0
+ MODE_SCALE: 1.29
+ DIFFUSION_MODEL:
+ NAME: FluxMR
+ PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@flux1-dev.safetensors
+ IN_CHANNELS: 64
+ OUT_CHANNELS: 64
+ HIDDEN_SIZE: 3072
+ NUM_HEADS: 24
+ AXES_DIM: [ 16, 56, 56 ]
+ THETA: 10000
+ VEC_IN_DIM: 768
+ GUIDANCE_EMBED: True
+ CONTEXT_IN_DIM: 4096
+ MLP_RATIO: 4.0
+ QKV_BIAS: True
+ DEPTH: 19
+ DEPTH_SINGLE_BLOCKS: 38
+ USE_GRAD_CHECKPOINT: True
+ ATTN_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKLFlux
+ EMBED_DIM: 16
+ PRETRAINED_MODEL: ms://AI-ModelScope/FLUX.1-dev@ae.safetensors
+ IGNORE_KEYS: [ ]
+ BATCH_SIZE: 8
+ USE_CONV: False
+ SCALE_FACTOR: 0.3611
+ SHIFT_FACTOR: 0.1159
+ #
+ ENCODER:
+ NAME: Encoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5PlusClipFluxEmbedder
+ T5_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: T5EncoderModel
+ MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder_2/
+ HF_TOKENIZER_CLS: T5Tokenizer
+ TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer_2/
+ MAX_LENGTH: 512
+ OUTPUT_KEY: last_hidden_state
+ D_TYPE: bfloat16
+ BATCH_INFER: False
+ CLEAN: whitespace
+ CLIP_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: CLIPTextModel
+ MODEL_PATH: ms://AI-ModelScope/FLUX.1-dev@text_encoder/
+ HF_TOKENIZER_CLS: CLIPTokenizer
+ TOKENIZER_PATH: ms://AI-ModelScope/FLUX.1-dev@tokenizer/
+ MAX_LENGTH: 77
+ OUTPUT_KEY: pooler_output
+ D_TYPE: bfloat16
+ BATCH_INFER: True
+ CLEAN: whitespace
+
+REFINER_MODEL_LOCAL:
+ NAME: ""
+ IS_DEFAULT: False
+ DEFAULT_PARAS:
+ PARAS:
+ RESOLUTIONS: [ [ 1024, 1024 ] ]
+ INPUT:
+ INPUT_IMAGE:
+ INPUT_MASK:
+ TASK:
+ PROMPT: ""
+ NEGATIVE_PROMPT: ""
+ OUTPUT_HEIGHT: 1024
+ OUTPUT_WIDTH: 1024
+ SAMPLER: flow_euler
+ SAMPLE_STEPS: 50
+ GUIDE_SCALE: 3.5
+ GUIDE_RESCALE: 0.5
+ OUTPUT:
+ LATENT:
+ IMAGES:
+ SEED:
+ MODULES_PARAS:
+ FIRST_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "IMAGE" ]
+ - NAME: decode
+ DTYPE: bfloat16
+ INPUT: [ "LATENT" ]
+ PARAS:
+ SCALE_FACTOR: 1.5305
+ SHIFT_FACTOR: 0.0609
+ SIZE_FACTOR: 8
+ DIFFUSION_MODEL:
+ FUNCTION:
+ - NAME: forward
+ DTYPE: bfloat16
+ INPUT: [ "SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE" ]
+ COND_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "PROMPT" ]
+
+ MODEL:
+ DIFFUSION:
+ NAME: DiffusionFluxRF
+ PREDICTION_TYPE: raw
+ NOISE_SCHEDULER:
+ NAME: FlowMatchSigmaScheduler
+ WEIGHTING_SCHEME: logit_normal
+ SHIFT: 3.0
+ LOGIT_MEAN: 0.0
+ LOGIT_STD: 1.0
+ MODE_SCALE: 1.29
+ DIFFUSION_MODEL:
+ NAME: FluxMR
+ PRETRAINED_MODEL: models/scepter/FLUX.1-dev/flux1-dev.safetensors
+ IN_CHANNELS: 64
+ OUT_CHANNELS: 64
+ HIDDEN_SIZE: 3072
+ NUM_HEADS: 24
+ AXES_DIM: [ 16, 56, 56 ]
+ THETA: 10000
+ VEC_IN_DIM: 768
+ GUIDANCE_EMBED: True
+ CONTEXT_IN_DIM: 4096
+ MLP_RATIO: 4.0
+ QKV_BIAS: True
+ DEPTH: 19
+ DEPTH_SINGLE_BLOCKS: 38
+ USE_GRAD_CHECKPOINT: True
+ ATTN_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKLFlux
+ EMBED_DIM: 16
+ PRETRAINED_MODEL: models/scepter/FLUX.1-dev/ae.safetensors
+ IGNORE_KEYS: [ ]
+ BATCH_SIZE: 8
+ USE_CONV: False
+ SCALE_FACTOR: 0.3611
+ SHIFT_FACTOR: 0.1159
+ #
+ ENCODER:
+ NAME: Encoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5PlusClipFluxEmbedder
+ T5_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: T5EncoderModel
+ MODEL_PATH: models/scepter/FLUX.1-dev/text_encoder_2/
+ HF_TOKENIZER_CLS: T5Tokenizer
+ TOKENIZER_PATH: models/scepter/FLUX.1-dev/FLUX.1-dev/tokenizer_2/
+ MAX_LENGTH: 512
+ OUTPUT_KEY: last_hidden_state
+ D_TYPE: bfloat16
+ BATCH_INFER: False
+ CLEAN: whitespace
+ CLIP_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: CLIPTextModel
+ MODEL_PATH: models/scepter/FLUX.1-dev/FLUX.1-dev/text_encoder/
+ HF_TOKENIZER_CLS: CLIPTokenizer
+ TOKENIZER_PATH: models/scepter/FLUX.1-dev/FLUX.1-dev/tokenizer/
+ MAX_LENGTH: 77
+ OUTPUT_KEY: pooler_output
+ D_TYPE: bfloat16
+ BATCH_INFER: True
+ CLEAN: whitespace
+
+REFINER_MODEL_HF:
+ NAME: ""
+ IS_DEFAULT: False
+ DEFAULT_PARAS:
+ PARAS:
+ RESOLUTIONS: [ [ 1024, 1024 ] ]
+ INPUT:
+ INPUT_IMAGE:
+ INPUT_MASK:
+ TASK:
+ PROMPT: ""
+ NEGATIVE_PROMPT: ""
+ OUTPUT_HEIGHT: 1024
+ OUTPUT_WIDTH: 1024
+ SAMPLER: flow_euler
+ SAMPLE_STEPS: 50
+ GUIDE_SCALE: 3.5
+ GUIDE_RESCALE: 0.5
+ OUTPUT:
+ LATENT:
+ IMAGES:
+ SEED:
+ MODULES_PARAS:
+ FIRST_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "IMAGE" ]
+ - NAME: decode
+ DTYPE: bfloat16
+ INPUT: [ "LATENT" ]
+ PARAS:
+ SCALE_FACTOR: 1.5305
+ SHIFT_FACTOR: 0.0609
+ SIZE_FACTOR: 8
+ DIFFUSION_MODEL:
+ FUNCTION:
+ - NAME: forward
+ DTYPE: bfloat16
+ INPUT: [ "SAMPLE_STEPS", "SAMPLE", "GUIDE_SCALE" ]
+ COND_STAGE_MODEL:
+ FUNCTION:
+ - NAME: encode
+ DTYPE: bfloat16
+ INPUT: [ "PROMPT" ]
+
+ MODEL:
+ DIFFUSION:
+ NAME: DiffusionFluxRF
+ PREDICTION_TYPE: raw
+ NOISE_SCHEDULER:
+ NAME: FlowMatchSigmaScheduler
+ WEIGHTING_SCHEME: logit_normal
+ SHIFT: 3.0
+ LOGIT_MEAN: 0.0
+ LOGIT_STD: 1.0
+ MODE_SCALE: 1.29
+ DIFFUSION_MODEL:
+ NAME: FluxMR
+ PRETRAINED_MODEL: hf://black-forest-labs/FLUX.1-dev@flux1-dev.safetensors
+ IN_CHANNELS: 64
+ OUT_CHANNELS: 64
+ HIDDEN_SIZE: 3072
+ NUM_HEADS: 24
+ AXES_DIM: [ 16, 56, 56 ]
+ THETA: 10000
+ VEC_IN_DIM: 768
+ GUIDANCE_EMBED: True
+ CONTEXT_IN_DIM: 4096
+ MLP_RATIO: 4.0
+ QKV_BIAS: True
+ DEPTH: 19
+ DEPTH_SINGLE_BLOCKS: 38
+ USE_GRAD_CHECKPOINT: True
+ ATTN_BACKEND: flash_attn
+ #
+ FIRST_STAGE_MODEL:
+ NAME: AutoencoderKLFlux
+ EMBED_DIM: 16
+ PRETRAINED_MODEL: hf://black-forest-labs/FLUX.1-dev@ae.safetensors
+ IGNORE_KEYS: [ ]
+ BATCH_SIZE: 8
+ USE_CONV: False
+ SCALE_FACTOR: 0.3611
+ SHIFT_FACTOR: 0.1159
+ #
+ ENCODER:
+ NAME: Encoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DOUBLE_Z: True
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ #
+ DECODER:
+ NAME: Decoder
+ USE_CHECKPOINT: False
+ CH: 128
+ OUT_CH: 3
+ NUM_RES_BLOCKS: 2
+ IN_CHANNELS: 3
+ ATTN_RESOLUTIONS: [ ]
+ CH_MULT: [ 1, 2, 4, 4 ]
+ Z_CHANNELS: 16
+ DROPOUT: 0.0
+ RESAMP_WITH_CONV: True
+ GIVE_PRE_END: False
+ TANH_OUT: False
+ #
+ COND_STAGE_MODEL:
+ NAME: T5PlusClipFluxEmbedder
+ T5_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: T5EncoderModel
+ MODEL_PATH: hf://black-forest-labs/FLUX.1-dev@text_encoder_2/
+ HF_TOKENIZER_CLS: T5Tokenizer
+ TOKENIZER_PATH: hf://black-forest-labs/FLUX.1-dev@tokenizer_2/
+ MAX_LENGTH: 512
+ OUTPUT_KEY: last_hidden_state
+ D_TYPE: bfloat16
+ BATCH_INFER: False
+ CLEAN: whitespace
+ CLIP_MODEL:
+ NAME: HFEmbedder
+ HF_MODEL_CLS: CLIPTextModel
+ MODEL_PATH: hf://black-forest-labs/FLUX.1-dev@text_encoder/
+ HF_TOKENIZER_CLS: CLIPTokenizer
+ TOKENIZER_PATH: hf://black-forest-labs/FLUX.1-dev@tokenizer/
+ MAX_LENGTH: 77
+ OUTPUT_KEY: pooler_output
+ D_TYPE: bfloat16
+ BATCH_INFER: True
+ CLEAN: whitespace
\ No newline at end of file
diff --git a/scepter/workflow/config/ace_0.6b_512_pro.yaml b/scepter/workflow/config/ace_0.6b_512_pro.yaml
index e7c29f7..c5bc98e 100644
--- a/scepter/workflow/config/ace_0.6b_512_pro.yaml
+++ b/scepter/workflow/config/ace_0.6b_512_pro.yaml
@@ -39,7 +39,7 @@ DEFAULT_PARAS:
#
COND_STAGE_MODEL:
FUNCTION:
- - NAME: encode_list
+ - NAME: encode_list_of_list
DTYPE: bfloat16
INPUT: ["PROMPT"]
#
diff --git a/scepter/workflow/config/flux1.0_dev_pro.yaml b/scepter/workflow/config/flux1.0_dev_pro.yaml
index 3e1f5d9..8e2d0dc 100644
--- a/scepter/workflow/config/flux1.0_dev_pro.yaml
+++ b/scepter/workflow/config/flux1.0_dev_pro.yaml
@@ -13,8 +13,8 @@ DEFAULT_PARAS:
VISIBLE: False
PROMPT_PREFIX: ""
SAMPLE:
- VALUES: ["flow_eluer"]
- DEFAULT: "flow_eluer"
+ VALUES: ["flow_euler"]
+ DEFAULT: "flow_euler"
SAMPLE_STEPS: 50
GUIDE_SCALE: 3.5
GUIDE_RESCALE:
diff --git a/scepter/workflow/config/flux1.0_schnell_pro.yaml b/scepter/workflow/config/flux1.0_schnell_pro.yaml
index b86a958..175b8bf 100644
--- a/scepter/workflow/config/flux1.0_schnell_pro.yaml
+++ b/scepter/workflow/config/flux1.0_schnell_pro.yaml
@@ -13,8 +13,8 @@ DEFAULT_PARAS:
VISIBLE: False
PROMPT_PREFIX: ""
SAMPLE:
- VALUES: ["flow_eluer"]
- DEFAULT: "flow_eluer"
+ VALUES: ["flow_euler"]
+ DEFAULT: "flow_euler"
SAMPLE_STEPS: 4
GUIDE_SCALE: 3.5
GUIDE_RESCALE:
diff --git a/scepter/workflow/config/scepter_workflow.yaml b/scepter/workflow/config/scepter_workflow.yaml
index 1afc620..fbb2fa5 100644
--- a/scepter/workflow/config/scepter_workflow.yaml
+++ b/scepter/workflow/config/scepter_workflow.yaml
@@ -59,6 +59,18 @@ BASE_MODELS:
FIRST_STAGE_MODEL: ACE_0.6B_512_AutoencoderKL
COND_STAGE_MODEL: ACE_0.6B_512_T5EmbedderHF
CONFIG: config/ace_0.6b_512_pro.yaml
+ -
+ NAME: ACE_0.6B_1024
+ DIFFUSION_MODEL: ACE_0.6B_1024_ACE
+ FIRST_STAGE_MODEL: ACE_0.6B_1024_AutoencoderKL
+ COND_STAGE_MODEL: ACE_0.6B_1024_T5EmbedderHF
+ CONFIG: config/ace_0.6b_1024_pro.yaml
+ -
+ NAME: ACE_0.6B_1024_REFINER
+ DIFFUSION_MODEL: ACE_0.6B_1024_REFINER_ACE
+ FIRST_STAGE_MODEL: ACE_0.6B_1024_REFINER_AutoencoderKL
+ COND_STAGE_MODEL: ACE_0.6B_1024_REFINER_T5EmbedderHF
+ CONFIG: config/ace_0.6b_1024_refiner_pro.yaml
MODEL_SOURCE:
- "ModelScope"
@@ -83,7 +95,7 @@ BASE_PARAMETERS:
- "dpmpp_2m_karras"
- "dpmpp_sde_karras"
- "dpmpp_2m_sde_karras"
- - "flow_eluer"
+ - "flow_euler"
DISCRETIZATION:
- "trailing"
diff --git a/scepter/workflow/model_node.py b/scepter/workflow/model_node.py
index 0e1b3f9..99c6555 100644
--- a/scepter/workflow/model_node.py
+++ b/scepter/workflow/model_node.py
@@ -94,10 +94,14 @@ class ModelNode:
elif source == 'Local':
cfg_new = copy.deepcopy(cfg)
cfg_new.MODEL = cfg_new.MODEL_LOCAL
+ if hasattr(cfg_new, 'EFINER_MODEL_LOCAL'):
+ cfg_new.REFINER_MODEL = cfg_new.REFINER_MODEL_LOCAL
return cfg_new
elif source == 'HuggingFace':
cfg_new = copy.deepcopy(cfg)
cfg_new.MODEL = cfg_new.MODEL_HF
+ if hasattr(cfg_new, 'EFINER_MODEL_HF'):
+ cfg_new.REFINER_MODEL = cfg_new.REFINER_MODEL_HF
return cfg_new
else:
raise NotImplementedError(f'Unknown model source: {source}')
diff --git a/tests/modules/test_diffusion_inference.py b/tests/modules/test_diffusion_inference.py
index 7dd035c..1d2346d 100644
--- a/tests/modules/test_diffusion_inference.py
+++ b/tests/modules/test_diffusion_inference.py
@@ -4,6 +4,7 @@
import os
import unittest
+import imageio
import numpy as np
import torchvision.transforms as TT
import torchvision.transforms.functional as TF
@@ -13,13 +14,15 @@ from torchvision.utils import save_image
from scepter.modules.annotator.registry import ANNOTATORS
from scepter.modules.inference.ace_inference import ACEInference
from scepter.modules.inference.diffusion_inference import DiffusionInference
-from scepter.modules.inference.flux_inference import FluxInference
from scepter.modules.inference.sd3_inference import SD3Inference
+from scepter.modules.inference.flux_inference import FluxInference
from scepter.modules.inference.stylebooth_inference import StyleboothInference
+from scepter.modules.inference.cogvideox_inference import CogVideoXInference
from scepter.modules.utils.config import Config
from scepter.modules.utils.distribute import we
from scepter.modules.utils.file_system import FS
from scepter.modules.utils.logger import get_logger
+from torchvision.utils import save_image
class DiffusionInferenceTest(unittest.TestCase):
@@ -235,10 +238,12 @@ class DiffusionInferenceTest(unittest.TestCase):
cfg = Config(cfg_file=config_file)
diff_infer = SD3Inference(logger=self.logger)
diff_infer.init_from_cfg(cfg)
- output = diff_infer({
- 'prompt': 'a cat holds a blackboard that writes "hello world"',
+ input_params = {
'seed': 2024
- })
+ }
+ output = diff_infer({
+ 'prompt': 'a cat holds a blackboard that writes "hello world"'
+ }, **input_params)
save_path = os.path.join(self.tmp_dir, 'sd3_cat.png')
save_image(output['images'], save_path)
print(save_path)
@@ -249,12 +254,17 @@ class DiffusionInferenceTest(unittest.TestCase):
cfg = Config(cfg_file=config_file)
diff_infer = FluxInference(logger=self.logger)
diff_infer.init_from_cfg(cfg)
- output = diff_infer({'prompt': '1 girl', 'seed': 2024})
+ input_params = {
+ 'seed': 2024
+ }
+ output = diff_infer({
+ 'prompt': '1 girl'
+ }, **input_params)
save_path = os.path.join(self.tmp_dir, 'flux_dev_1girl.png')
save_image(output['images'], save_path)
print(save_path)
- # @unittest.skip('')
+ @unittest.skip('')
def test_ace(self):
config_file = 'scepter/methods/studio/chatbot/models/ace_0.6b_512.yaml'
cfg = Config(cfg_file=config_file)
@@ -266,5 +276,26 @@ class DiffusionInferenceTest(unittest.TestCase):
print(save_path)
+ # @unittest.skip('')
+ def test_cogvideox_2b(self):
+ config_file = 'scepter/methods/studio/inference/dit/cogvideox_2b_pro.yaml'
+ cfg = Config(cfg_file=config_file)
+ diff_infer = CogVideoXInference(logger=self.logger)
+ diff_infer.init_from_cfg(cfg)
+ input_params = {
+ 'seed': 42
+ }
+ output = diff_infer({
+ 'prompt': 'A girl riding a bike.'
+ }, **input_params)
+ frames = (output['videos'][0].permute(1, 2, 3, 0).cpu().numpy() * 255).astype(np.uint8)
+ save_path = os.path.join(self.tmp_dir, 'cogvideox_2b_girlbike.mp4')
+ writer = imageio.get_writer(save_path, fps=8)
+ for frame in frames:
+ writer.append_data(np.array(frame))
+ writer.close()
+ print(save_path)
+
+
if __name__ == '__main__':
unittest.main()