| 317 | } |
| 318 | |
| 319 | bool StableDiffusion::run(const std::string prompt, const std::string imagePath, int iterNum, int randomSeed, std::function<void(int)> progressCallback) { |
| 320 | AUTOTIME; |
| 321 | mEts.clear(); |
| 322 | |
| 323 | if(iterNum > 50) { |
| 324 | iterNum = 50; |
| 325 | MNN_PRINT("too much number of iterations, iterations will be set to 50.\n"); |
| 326 | } |
| 327 | if(iterNum < 1) { |
| 328 | iterNum = 10; |
| 329 | MNN_PRINT("illegal number of iterations, iterations will be set to 10.\n"); |
| 330 | } |
| 331 | mTimeSteps.resize(iterNum); |
| 332 | int step = 1000 / iterNum; |
| 333 | for(int i = iterNum - 1; i >= 0; i--) { |
| 334 | mTimeSteps[i] = 1 + (iterNum - 1 - i) * step; |
| 335 | } |
| 336 | |
| 337 | auto ids = mTokenizer->encode(prompt, mMaxTextLen); |
| 338 | |
| 339 | auto text_embeddings = text_encoder(ids); |
| 340 | |
| 341 | if (progressCallback) { |
| 342 | progressCallback(1 * 100 / (iterNum + 3)); // percent |
| 343 | } |
| 344 | auto latent = unet(text_embeddings, iterNum, randomSeed, progressCallback); |
| 345 | |
| 346 | auto image = vae_decoder(latent); |
| 347 | bool res = imwrite(imagePath, image); |
| 348 | if (res) { |
| 349 | MNN_PRINT("SUCCESS! write generated image to %s\n", imagePath.c_str()); |
| 350 | } |
| 351 | |
| 352 | if(mMemoryMode != 1) { |
| 353 | mModules[2].reset(); |
| 354 | } |
| 355 | |
| 356 | if (progressCallback) { |
| 357 | progressCallback(100); // percent |
| 358 | } |
| 359 | return true; |
| 360 | } |
| 361 | |
| 362 | |
| 363 | // 统一的生成接口实现 |