fred-dev/comfy_ui_ali

1import comfy.options2comfy.options.enable_args_parsing()3 4import os5import importlib.util6import folder_paths7import time8from comfy.cli_args import args9from app.logger import setup_logger10import itertools11import utils.extra_config12import logging13 14if __name__ == "__main__":15    #NOTE: These do not do anything on core ComfyUI which should already have no communication with the internet, they are for custom nodes.16 17    os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'18    os.environ['DO_NOT_TRACK'] = '1'19 20    # make the directories21    os.makedirs("/data/models/checkpoints/", exist_ok=True)22    os.makedirs("/data/models/clip/", exist_ok=True)23    os.makedirs("/data/models/clip_vision/", exist_ok=True)24    os.makedirs("/data/models/configs/", exist_ok=True)25    os.makedirs("/data/models/controlnet/", exist_ok=True)26    os.makedirs("/data/models/diffusion_models/", exist_ok=True)27    os.makedirs("/data/models/unet/", exist_ok=True)28    os.makedirs("/data/models/embeddings/", exist_ok=True)29    os.makedirs("/data/models/loras/", exist_ok=True)30    os.makedirs("/data/models/upscale_models/", exist_ok=True)31    os.makedirs("/data/models/vae/", exist_ok=True)32    os.makedirs("/data/output/", exist_ok=True)33    os.makedirs("/data/temp/", exist_ok=True)34    os.makedirs("/data/input/", exist_ok=True)35    os.makedirs("/data/user/", exist_ok=True)36 37    #lets print the contents of the directories and subdirectories from the root directory38    for root, dirs, files in os.walk("/data"):39        print("Printing contents of the directory: ", root)40        for name in files:41            print(os.path.join(root, name))42        for name in dirs:43            print(os.path.join(root, name))44 45    46    #lets print the contents of the directories and subdirectories from the root directory47    for root, dirs, files in os.walk("/data"):48        print("Printing contents of the directory: ", root)49        for name in files:50            print(os.path.join(root, name))51        for name in dirs:52            print(os.path.join(root, name))53 54 55 56setup_logger(log_level=args.verbose, use_stdout=args.log_stdout)57 58def apply_custom_paths():59    # extra model paths60    extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml")61    if os.path.isfile(extra_model_paths_config_path):62        utils.extra_config.load_extra_path_config(extra_model_paths_config_path)63        print(f"Loaded extra model paths from: {extra_model_paths_config_path}")64 65    if args.extra_model_paths_config:66        for config_path in itertools.chain(*args.extra_model_paths_config):67            utils.extra_config.load_extra_path_config(config_path)68 69    # --output-directory, --input-directory, --user-directory70    if args.output_directory:71        output_dir = os.path.abspath(args.output_directory)72        logging.info(f"Setting output directory to: {output_dir}")73        folder_paths.set_output_directory(output_dir)74 75    # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes76    folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints"))77    folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip"))78    folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae"))79    folder_paths.add_model_folder_path("diffusion_models",80                                       os.path.join(folder_paths.get_output_directory(), "diffusion_models"))81    folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras"))82 83    if args.input_directory:84        input_dir = os.path.abspath(args.input_directory)85        logging.info(f"Setting input directory to: {input_dir}")86        folder_paths.set_input_directory(input_dir)87 88    if args.user_directory:89        user_dir = os.path.abspath(args.user_directory)90        logging.info(f"Setting user directory to: {user_dir}")91        folder_paths.set_user_directory(user_dir)92 93 94def execute_prestartup_script():95    def execute_script(script_path):96        module_name = os.path.splitext(script_path)[0]97        try:98            spec = importlib.util.spec_from_file_location(module_name, script_path)99            module = importlib.util.module_from_spec(spec)100            spec.loader.exec_module(module)101            return True102        except Exception as e:103            logging.error(f"Failed to execute startup-script: {script_path} / {e}")104        return False105 106    if args.disable_all_custom_nodes:107        return108 109    node_paths = folder_paths.get_folder_paths("custom_nodes")110    for custom_node_path in node_paths:111        possible_modules = os.listdir(custom_node_path)112        node_prestartup_times = []113 114        for possible_module in possible_modules:115            module_path = os.path.join(custom_node_path, possible_module)116            if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__":117                continue118 119            script_path = os.path.join(module_path, "prestartup_script.py")120            if os.path.exists(script_path):121                time_before = time.perf_counter()122                success = execute_script(script_path)123                node_prestartup_times.append((time.perf_counter() - time_before, module_path, success))124    if len(node_prestartup_times) > 0:125        logging.info("\nPrestartup times for custom nodes:")126        for n in sorted(node_prestartup_times):127            if n[2]:128                import_message = ""129            else:130                import_message = " (PRESTARTUP FAILED)"131            logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1]))132        logging.info("")133 134apply_custom_paths()135execute_prestartup_script()136 137 138# Main code139import asyncio140import shutil141import threading142import gc143 144 145if os.name == "nt":146    logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())147 148if __name__ == "__main__":149    if args.cuda_device is not None:150        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)151        os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device)152        logging.info("Set cuda device to: {}".format(args.cuda_device))153 154    if args.oneapi_device_selector is not None:155        os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector156        logging.info("Set oneapi device selector to: {}".format(args.oneapi_device_selector))157 158    if args.deterministic:159        if 'CUBLAS_WORKSPACE_CONFIG' not in os.environ:160            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"161 162    import cuda_malloc163 164if args.windows_standalone_build:165    try:166        from fix_torch import fix_pytorch_libomp167        fix_pytorch_libomp()168    except:169        pass170 171import comfy.utils172 173import execution174import server175from server import BinaryEventTypes176import nodes177import comfy.model_management178import comfyui_version179import app.logger180 181 182def cuda_malloc_warning():183    device = comfy.model_management.get_torch_device()184    device_name = comfy.model_management.get_torch_device_name(device)185    cuda_malloc_warning = False186    if "cudaMallocAsync" in device_name:187        for b in cuda_malloc.blacklist:188            if b in device_name:189                cuda_malloc_warning = True190        if cuda_malloc_warning:191            logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")192 193 194def prompt_worker(q, server_instance):195    current_time: float = 0.0196    e = execution.PromptExecutor(server_instance, lru_size=args.cache_lru)197    last_gc_collect = 0198    need_gc = False199    gc_collect_interval = 10.0200 201    while True:202        timeout = 1000.0203        if need_gc:204            timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)205 206        queue_item = q.get(timeout=timeout)207        if queue_item is not None:208            item, item_id = queue_item209            execution_start_time = time.perf_counter()210            prompt_id = item[1]211            server_instance.last_prompt_id = prompt_id212 213            e.execute(item[2], prompt_id, item[3], item[4])214            need_gc = True215            q.task_done(item_id,216                        e.history_result,217                        status=execution.PromptQueue.ExecutionStatus(218                            status_str='success' if e.success else 'error',219                            completed=e.success,220                            messages=e.status_messages))221            if server_instance.client_id is not None:222                server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)223 224            current_time = time.perf_counter()225            execution_time = current_time - execution_start_time226            logging.info("Prompt executed in {:.2f} seconds".format(execution_time))227 228        flags = q.get_flags()229        free_memory = flags.get("free_memory", False)230 231        if flags.get("unload_models", free_memory):232            comfy.model_management.unload_all_models()233            need_gc = True234            last_gc_collect = 0235 236        if free_memory:237            e.reset()238            need_gc = True239            last_gc_collect = 0240 241        if need_gc:242            current_time = time.perf_counter()243            if (current_time - last_gc_collect) > gc_collect_interval:244                gc.collect()245                comfy.model_management.soft_empty_cache()246                last_gc_collect = current_time247                need_gc = False248 249 250async def run(server_instance, address='0.0.0.0', port=7860, verbose=True, call_on_start=None):251    addresses = []252    for addr in address.split(","):253        addresses.append((addr, port))254    await asyncio.gather(255        server_instance.start_multi_address(addresses, call_on_start, verbose), server_instance.publish_loop()256    )257 258 259def hijack_progress(server_instance):260    def hook(value, total, preview_image):261        comfy.model_management.throw_exception_if_processing_interrupted()262        progress = {"value": value, "max": total, "prompt_id": server_instance.last_prompt_id, "node": server_instance.last_node_id}263 264        server_instance.send_sync("progress", progress, server_instance.client_id)265        if preview_image is not None:266            server_instance.send_sync(BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, preview_image, server_instance.client_id)267 268    comfy.utils.set_progress_bar_global_hook(hook)269 270 271def cleanup_temp():272    temp_dir = folder_paths.get_temp_directory()273    if os.path.exists(temp_dir):274        shutil.rmtree(temp_dir, ignore_errors=True)275 276 277def start_comfyui(asyncio_loop=None):278    """279    Starts the ComfyUI server using the provided asyncio event loop or creates a new one.280    Returns the event loop, server instance, and a function to start the server asynchronously.281    """282    if args.temp_directory:283        temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")284        logging.info(f"Setting temp directory to: {temp_dir}")285        folder_paths.set_temp_directory(temp_dir)286    cleanup_temp()287 288    if args.windows_standalone_build:289        try:290            import new_updater291            new_updater.update_windows_updater()292        except:293            pass294 295    if not asyncio_loop:296        asyncio_loop = asyncio.new_event_loop()297        asyncio.set_event_loop(asyncio_loop)298    prompt_server = server.PromptServer(asyncio_loop)299    q = execution.PromptQueue(prompt_server)300 301    nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)302 303    cuda_malloc_warning()304 305    prompt_server.add_routes()306    hijack_progress(prompt_server)307 308    threading.Thread(target=prompt_worker, daemon=True, args=(q, prompt_server,)).start()309 310    if args.quick_test_for_ci:311        exit(0)312 313    os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)314    call_on_start = None315    if args.auto_launch:316        def startup_server(scheme, address, port):317            import webbrowser318            if os.name == 'nt' and address == '0.0.0.0':319                address = '0.0.0.0'320            if ':' in address:321                address = "[{}]".format(address)322            webbrowser.open(f"{scheme}://{address}:{port}")323        call_on_start = startup_server324 325    async def start_all():326        await prompt_server.setup()327        await run(prompt_server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start)328 329    # Returning these so that other code can integrate with the ComfyUI loop and server330    return asyncio_loop, prompt_server, start_all331 332 333if __name__ == "__main__":334    # Running directly, just start ComfyUI.335    logging.info("ComfyUI version: {}".format(comfyui_version.__version__))336 337    event_loop, _, start_all_func = start_comfyui()338    try:339        x = start_all_func()340        app.logger.print_startup_warnings()341        event_loop.run_until_complete(x)342    except KeyboardInterrupt:343        logging.info("\nStopped server")344 345    cleanup_temp()346