def _load_models_to_neuron()

in optimum/neuron/modeling_diffusion.py [0:0]

• 37 lines of code
• 19 McCabe index (conditional complexity)

        def _load_models_to_neuron(submodels, models_on_both_cores=None, models_on_a_single_core=None):
            # loading models to both cores, eg. unet, transformer.
            if models_on_both_cores:
                for model_name in models_on_both_cores:
                    submodel_paths = submodels[model_name]
                    # for the case of multiple controlnets the path could be a list
                    if not isinstance(submodel_paths, list):
                        submodel_paths = [submodel_paths]
                    submodels_list = []
                    for submodel_path in submodel_paths:
                        if submodel_path is not None and submodel_path.is_file():
                            submodel = NeuronTracedModel.load_model(
                                submodel_path, to_neuron=False
                            )  # No need to load to neuron manually when dp
                            submodel = torch_neuronx.DataParallel(
                                submodel,
                                [0, 1],
                                set_dynamic_batching=dynamic_batch_size,
                            )
                            submodels_list.append(submodel)
                    if submodels_list:
                        submodels[model_name] = submodels_list if len(submodels_list) > 1 else submodels_list[0]
                    else:
                        submodels[model_name] = None
            # loading models to a single core, eg. text encoders, vae.
            if models_on_a_single_core:
                for model_name in models_on_a_single_core:
                    submodel_paths = submodels[model_name]
                    # for the case of multiple controlnets the path could be a list
                    if not isinstance(submodel_paths, list):
                        submodel_paths = [submodel_paths]
                    submodels_list = []
                    for submodel_path in submodel_paths:
                        if submodel_path is not None and submodel_path.is_file():
                            submodel = NeuronTracedModel.load_model(submodel_path, to_neuron=to_neuron)
                            submodels_list.append(submodel)
                    if submodels_list:
                        submodels[model_name] = submodels_list if len(submodels_list) > 1 else submodels_list[0]
                    else:
                        submodels[model_name] = None
            return submodels