File size: 11,506 Bytes
a1da63c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
from typing import List, Optional, Tuple

import gradio

from facefusion import logger, state_manager, wording
from facefusion.args import collect_step_args
from facefusion.common_helper import get_first, get_last
from facefusion.filesystem import is_directory
from facefusion.jobs import job_manager
from facefusion.typing import UiWorkflow
from facefusion.uis import choices as uis_choices
from facefusion.uis.core import get_ui_component
from facefusion.uis.typing import JobManagerAction
from facefusion.uis.ui_helper import convert_int_none, convert_str_none, suggest_output_path

JOB_MANAGER_WRAPPER : Optional[gradio.Column] = None
JOB_MANAGER_JOB_ACTION_DROPDOWN : Optional[gradio.Dropdown] = None
JOB_MANAGER_JOB_ID_TEXTBOX : Optional[gradio.Textbox] = None
JOB_MANAGER_JOB_ID_DROPDOWN : Optional[gradio.Dropdown] = None
JOB_MANAGER_STEP_INDEX_DROPDOWN : Optional[gradio.Dropdown] = None
JOB_MANAGER_APPLY_BUTTON : Optional[gradio.Button] = None


def render() -> None:
	global JOB_MANAGER_WRAPPER
	global JOB_MANAGER_JOB_ACTION_DROPDOWN
	global JOB_MANAGER_JOB_ID_TEXTBOX
	global JOB_MANAGER_JOB_ID_DROPDOWN
	global JOB_MANAGER_STEP_INDEX_DROPDOWN
	global JOB_MANAGER_APPLY_BUTTON

	if job_manager.init_jobs(state_manager.get_item('jobs_path')):
		is_job_manager = state_manager.get_item('ui_workflow') == 'job_manager'
		drafted_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ]

		with gradio.Column(visible = is_job_manager) as JOB_MANAGER_WRAPPER:
			JOB_MANAGER_JOB_ACTION_DROPDOWN = gradio.Dropdown(
				label = wording.get('uis.job_manager_job_action_dropdown'),
				choices = uis_choices.job_manager_actions,
				value = get_first(uis_choices.job_manager_actions)
			)
			JOB_MANAGER_JOB_ID_TEXTBOX = gradio.Textbox(
				label = wording.get('uis.job_manager_job_id_dropdown'),
				max_lines = 1,
				interactive = True
			)
			JOB_MANAGER_JOB_ID_DROPDOWN = gradio.Dropdown(
				label = wording.get('uis.job_manager_job_id_dropdown'),
				choices = drafted_job_ids,
				value = get_last(drafted_job_ids),
				interactive = True,
				visible = False
			)
			JOB_MANAGER_STEP_INDEX_DROPDOWN = gradio.Dropdown(
				label = wording.get('uis.job_manager_step_index_dropdown'),
				choices = [ 'none' ],
				value = 'none',
				interactive = True,
				visible = False
			)
			JOB_MANAGER_APPLY_BUTTON = gradio.Button(
				value = wording.get('uis.apply_button'),
				variant = 'primary',
				size = 'sm'
			)


def listen() -> None:
	JOB_MANAGER_JOB_ACTION_DROPDOWN.change(update, inputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_DROPDOWN ], outputs = [ JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ])
	JOB_MANAGER_JOB_ID_DROPDOWN.change(update_step_index, inputs = JOB_MANAGER_JOB_ID_DROPDOWN, outputs = JOB_MANAGER_STEP_INDEX_DROPDOWN)
	JOB_MANAGER_APPLY_BUTTON.click(apply, inputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ], outputs = [ JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ])

	ui_workflow_dropdown = get_ui_component('ui_workflow_dropdown')
	if ui_workflow_dropdown:
		ui_workflow_dropdown.change(remote_update, inputs = ui_workflow_dropdown, outputs = [ JOB_MANAGER_WRAPPER, JOB_MANAGER_JOB_ACTION_DROPDOWN, JOB_MANAGER_JOB_ID_TEXTBOX, JOB_MANAGER_JOB_ID_DROPDOWN, JOB_MANAGER_STEP_INDEX_DROPDOWN ])


def remote_update(ui_workflow : UiWorkflow) -> Tuple[gradio.Row, gradio.Dropdown, gradio.Textbox, gradio.Dropdown, gradio.Dropdown]:
	is_job_manager = ui_workflow == 'job_manager'
	return gradio.Row(visible = is_job_manager), gradio.Dropdown(value = get_first(uis_choices.job_manager_actions)), gradio.Textbox(value = None, visible = True), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False)


def apply(job_action : JobManagerAction, created_job_id : str, selected_job_id : str, selected_step_index : int) -> Tuple[gradio.Dropdown, gradio.Textbox, gradio.Dropdown, gradio.Dropdown]:
	created_job_id = convert_str_none(created_job_id)
	selected_job_id = convert_str_none(selected_job_id)
	selected_step_index = convert_int_none(selected_step_index)
	step_args = collect_step_args()
	output_path = step_args.get('output_path')

	if is_directory(step_args.get('output_path')):
		step_args['output_path'] = suggest_output_path(step_args.get('output_path'), state_manager.get_item('target_path'))
	if job_action == 'job-create':
		if created_job_id and job_manager.create_job(created_job_id):
			updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ]

			logger.info(wording.get('job_created').format(job_id = created_job_id), __name__.upper())
			return gradio.Dropdown(value = 'job-add-step'), gradio.Textbox(visible = False), gradio.Dropdown(value = created_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown()
		else:
			logger.error(wording.get('job_not_created').format(job_id = created_job_id), __name__.upper())
	if job_action == 'job-submit':
		if selected_job_id and job_manager.submit_job(selected_job_id):
			updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ]

			logger.info(wording.get('job_submitted').format(job_id = selected_job_id), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True), gradio.Dropdown()
		else:
			logger.error(wording.get('job_not_submitted').format(job_id = selected_job_id), __name__.upper())
	if job_action == 'job-delete':
		if selected_job_id and job_manager.delete_job(selected_job_id):
			updated_job_ids = job_manager.find_job_ids('drafted') + job_manager.find_job_ids('queued') + job_manager.find_job_ids('failed') + job_manager.find_job_ids('completed') or [ 'none' ]

			logger.info(wording.get('job_deleted').format(job_id = selected_job_id), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(value = get_last(updated_job_ids), choices = updated_job_ids, visible = True), gradio.Dropdown()
		else:
			logger.error(wording.get('job_not_deleted').format(job_id = selected_job_id), __name__.upper())
	if job_action == 'job-add-step':
		if selected_job_id and job_manager.add_step(selected_job_id, step_args):
			state_manager.set_item('output_path', output_path)
			logger.info(wording.get('job_step_added').format(job_id = selected_job_id), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(visible = False)
		else:
			state_manager.set_item('output_path', output_path)
			logger.error(wording.get('job_step_not_added').format(job_id = selected_job_id), __name__.upper())
	if job_action == 'job-remix-step':
		if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.remix_step(selected_job_id, selected_step_index, step_args):
			updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item]

			state_manager.set_item('output_path', output_path)
			logger.info(wording.get('job_remix_step_added').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True)
		else:
			state_manager.set_item('output_path', output_path)
			logger.error(wording.get('job_remix_step_not_added').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
	if job_action == 'job-insert-step':
		if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.insert_step(selected_job_id, selected_step_index, step_args):
			updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item]

			state_manager.set_item('output_path', output_path)
			logger.info(wording.get('job_step_inserted').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True)
		else:
			state_manager.set_item('output_path', output_path)
			logger.error(wording.get('job_step_not_inserted').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
	if job_action == 'job-remove-step':
		if selected_job_id and job_manager.has_step(selected_job_id, selected_step_index) and job_manager.remove_step(selected_job_id, selected_step_index):
			updated_step_choices = get_step_choices(selected_job_id) or [ 'none' ] #type:ignore[list-item]

			logger.info(wording.get('job_step_removed').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
			return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True)
		else:
			logger.error(wording.get('job_step_not_removed').format(job_id = selected_job_id, step_index = selected_step_index), __name__.upper())
	return gradio.Dropdown(), gradio.Textbox(), gradio.Dropdown(), gradio.Dropdown()


def get_step_choices(job_id : str) -> List[int]:
	steps = job_manager.get_steps(job_id)
	return [ index for index, _ in enumerate(steps) ]


def update(job_action : JobManagerAction, selected_job_id : str) -> Tuple[gradio.Textbox, gradio.Dropdown, gradio.Dropdown]:
	if job_action == 'job-create':
		return gradio.Textbox(value = None, visible = True), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False)
	if job_action == 'job-delete':
		updated_job_ids = job_manager.find_job_ids('drafted') + job_manager.find_job_ids('queued') + job_manager.find_job_ids('failed') + job_manager.find_job_ids('completed') or [ 'none' ]
		updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids)

		return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(visible = False)
	if job_action in [ 'job-submit', 'job-add-step' ]:
		updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ]
		updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids)

		return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(visible = False)
	if job_action in [ 'job-remix-step', 'job-insert-step', 'job-remove-step' ]:
		updated_job_ids = job_manager.find_job_ids('drafted') or [ 'none' ]
		updated_job_id = selected_job_id if selected_job_id in updated_job_ids else get_last(updated_job_ids)
		updated_step_choices = get_step_choices(updated_job_id) or [ 'none' ] #type:ignore[list-item]

		return gradio.Textbox(visible = False), gradio.Dropdown(value = updated_job_id, choices = updated_job_ids, visible = True), gradio.Dropdown(value = get_last(updated_step_choices), choices = updated_step_choices, visible = True)
	return gradio.Textbox(visible = False), gradio.Dropdown(visible = False), gradio.Dropdown(visible = False)


def update_step_index(job_id : str) -> gradio.Dropdown:
	step_choices = get_step_choices(job_id) or [ 'none' ] #type:ignore[list-item]
	return gradio.Dropdown(value = get_last(step_choices), choices = step_choices)