MCPcopy Create free account
hub / github.com/alibaba/MNN / load

Method load

transformers/diffusion/engine/src/stable_diffusion.cpp:42–160  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

40}
41
42bool StableDiffusion::load() {
43 AUTOTIME;
44#if !defined(MNN_DIFFUSION_WITH_LLM_TOKENIZER)
45 MNN_ERROR("Diffusion models require MNN_BUILD_LLM=ON so diffusion can load tokenizer.mtok\n");
46 return false;
47#endif
48 ScheduleConfig config;
49 BackendConfig backendConfig;
50 config.type = mBackendType;
51 if(config.type == MNN_FORWARD_CPU) {
52 config.numThread = 4;
53 } else if(config.type == MNN_FORWARD_OPENCL) {
54 config.mode = MNN_GPU_MEMORY_BUFFER | MNN_GPU_TUNING_FAST;
55 } else {
56 config.numThread = 1;
57 }
58 backendConfig.memory = BackendConfig::Memory_Low;
59 backendConfig.precision = BackendConfig::Precision_Low;
60 config.backendConfig = &backendConfig;
61
62 auto exe = ExecutorScope::Current();
63 exe->lazyEval = false;
64 exe->setGlobalExecutorConfig(config.type, backendConfig, config.numThread);
65
66 Module::Config module_config;
67 module_config.shapeMutable = false;
68 // module_config.rearrange = true;
69 runtime_manager_.reset(Executor::RuntimeManager::createRuntimeManager(config));
70
71 if (config.type == MNN_FORWARD_OPENCL) {
72 const char* cacheFileName = ".tempcache";
73 runtime_manager_->setCache(cacheFileName);
74 }
75 // need to consider memory
76 if(mMemoryMode == 0) {
77 runtime_manager_->setHint(Interpreter::WINOGRAD_MEMORY_LEVEL, 0);
78 } else if(mMemoryMode == 2) {
79 runtime_manager_->setHint(Interpreter::WINOGRAD_MEMORY_LEVEL, 1);
80 }
81 if(config.type == MNN_FORWARD_CPU) {
82 runtime_manager_->setHint(Interpreter::DYNAMIC_QUANT_OPTIONS, 2);
83 }
84 mLatentVar = _Input({1, 4, 64, 64}, NCHW, halide_type_of<float>());
85 mPromptVar = _Input({2, mMaxTextLen}, NCHW, halide_type_of<int>());
86 mTimestepVar = _Input({1}, NCHW, halide_type_of<int>());
87 mLatentVar->writeMap<int8_t>();
88 mPromptVar->writeMap<int8_t>();
89 mTimestepVar->writeMap<int8_t>();
90 mSampleVar = _Concat({mLatentVar, mLatentVar}, 0);
91
92 if(mMemoryMode > 0) {
93 MNN_PRINT("First time initilizing may cost a few seconds to create cachefile, please wait ...\n");
94 }
95
96 VARP text_embeddings;
97 mModules.resize(3);
98 mResizeCachePrepared.assign(3, false);
99 // load text_encoder model

Callers

nothing calls this directly

Calls 11

_InputFunction · 0.85
_ConcatFunction · 0.85
_ConvertFunction · 0.85
fixMethod · 0.80
loadFunction · 0.50
resetMethod · 0.45
setCacheMethod · 0.45
setHintMethod · 0.45
resizeMethod · 0.45
c_strMethod · 0.45

Tested by

no test coverage detected