diff --git a/paint/sources/neural_nodes/upscale_node.ts b/paint/sources/neural_nodes/upscale_node.ts index dc8e57c1..d233d850 100644 --- a/paint/sources/neural_nodes/upscale_node.ts +++ b/paint/sources/neural_nodes/upscale_node.ts @@ -1,5 +1,5 @@ -let upscale_result: gpu_texture_t = null; +let upscale_node_result: gpu_texture_t = null; function upscale_node_init() { array_push(nodes_material_neural, upscale_node_def); @@ -8,24 +8,75 @@ function upscale_node_init() { } function upscale_node_vector(node: ui_node_t, socket: ui_node_socket_t): string { - if (upscale_result == null) { + if (upscale_node_result == null) { return("float3(0.0, 0.0, 0.0)"); } let tex_name: string = parser_material_node_name(node); - map_set(data_cached_images, tex_name, upscale_result); + map_set(data_cached_images, tex_name, upscale_node_result); let tex: bind_tex_t = parser_material_make_bind_tex(tex_name, tex_name); let texstore: string = parser_material_texture_store(node, tex, tex_name, color_space_t.AUTO); return texstore + ".rgb"; } function upscale_node_button(node_id: i32) { - let node: ui_node_t = ui_get_node(ui_nodes_get_canvas(true).nodes, node_id); + let canvas: ui_node_canvas_t = ui_nodes_get_canvas(true); + let node: ui_node_t = ui_get_node(canvas.nodes, node_id); + + let models: string[] = ["RealESRGAN_X2", "RealESRGAN_X4"]; + let model: i32 = ui_combo(ui_handle(__ID__), models, tr("Model")); + + if (iron_exec_async_done == 0) { + ui_button("Cancel..."); + } + else if (ui_button("Run")) { + let inp: ui_node_socket_t = node.inputs[1]; + let from_node: ui_node_t = null; + for (let i: i32 = 0; i < canvas.links.length; ++i) { + let l: ui_node_link_t = canvas.links[i]; + if (l.to_id == inp.node_id) { + from_node = ui_get_node(canvas.nodes, l.from_id); + break; + } + } + + let input: gpu_texture_t = ui_nodes_get_node_preview_image(from_node); + if (input != null) { + let dir: string; + if (path_is_protected()) { + dir = iron_internal_save_path() + "models"; + } + else { + dir = iron_internal_files_location() + path_sep + "models"; + } + iron_write_png(dir + path_sep + "input.png", gpu_get_texture_pixels(input), input.width, input.height, 0); + let argv: string[] = [ + dir + "/sd", + "-M", "upscale", + "--upscale-model", dir + "/RealESRGAN_x4plus.pth", + "-i", dir + "/input.png", + "-o", dir + "/output.png", + null + ]; + iron_exec_async(argv[0], argv.buffer); + sys_notify_on_update(upscale_node_check_result, dir); + } + } +} + +function upscale_node_check_result(dir: string) { + if (iron_exec_async_done == 1) { + let file: string = dir + path_sep + "output.png"; + if (iron_file_exists(file)) { + upscale_node_result = iron_load_texture(file); + } + sys_remove_update(upscale_node_check_result); + } } let upscale_node_def: ui_node_t = { id: 0, name: _tr("Upscale"), - type: "upscale_node", + type: "NEURAL_UPSCALE", x: 0, y: 0, color: 0xff4982a0, @@ -69,7 +120,19 @@ let upscale_node_def: ui_node_t = { display: 0 } ], - buttons: [], + buttons: [ + { + name: "upscale_node_button", + type: "CUSTOM", + output: -1, + default_value: f32_array_create_x(0), + data: null, + min: 0.0, + max: 1.0, + precision: 100, + height: 2 + } + ], width: 0, flags: 0 }; diff --git a/paint/sources/ui_nodes.ts b/paint/sources/ui_nodes.ts index 4fac9137..9656ec87 100644 --- a/paint/sources/ui_nodes.ts +++ b/paint/sources/ui_nodes.ts @@ -1156,7 +1156,6 @@ function ui_nodes_get_node_preview_image(n: ui_node_t): gpu_texture_t { else if (n.type == "TEX_IMAGE" && parser_material_get_input_link(n.inputs[0]) == null) { let i: i32 = n.buttons[0].default_value[0]; let filepath: string = parser_material_enum_data(base_enum_texts(n.type)[i]); - let asset_index: i32 = -1; for (let i: i32 = 0; i < project_assets.length; ++i) { if (project_assets[i].file == filepath) { @@ -1164,11 +1163,16 @@ function ui_nodes_get_node_preview_image(n: ui_node_t): gpu_texture_t { break; } } - if (asset_index > -1) { img = project_get_image(project_assets[asset_index]); } } + else if (n.type == "NEURAL_TEXT_TO_IMAGE") { + img = text_to_image_node_result; + } + else if (n.type == "NEURAL_UPSCALE") { + img = upscale_node_result; + } else if (ui_nodes_canvas_type == canvas_type_t.MATERIAL) { img = any_imap_get(context_raw.node_preview_map, n.id); }