chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,158 @@
|
||||
//
|
||||
// sana_diffusion_demo.cpp
|
||||
//
|
||||
// Sana Diffusion 演示程序
|
||||
//
|
||||
// 展示如何使用Sana模型进行文生图和图像编辑：
|
||||
// 1. 使用Qwen3-0.6B LLM处理文本prompt
|
||||
// 2. 通过Connector和Projector桥接特征
|
||||
// 3. 使用DiT Transformer进行去噪生成
|
||||
// 4. VAE解码得到最终图像
|
||||
//
|
||||
|
||||
#include <iostream>
|
||||
#include "diffusion/diffusion.hpp"
|
||||
#include "diffusion/sana_llm.hpp"
|
||||
#define MNN_OPEN_TIME_TRACE
|
||||
#include <MNN/AutoTime.hpp>
|
||||
#include <MNN/expr/ExecutorScope.hpp>
|
||||
|
||||
using namespace MNN::DIFFUSION;
|
||||
using namespace MNN::Express;
|
||||
|
||||
int main(int argc, const char* argv[]) {
|
||||
if (argc < 3) {
|
||||
MNN_PRINT("=====================================================================================================================\n");
|
||||
MNN_PRINT("Sana Diffusion Demo - 基于Qwen3-0.6B的高效文生图模型\n");
|
||||
MNN_PRINT("=====================================================================================================================\n");
|
||||
MNN_PRINT("Usage: ./sana_diffusion_demo <resource_path> <mode> <prompt> [input_image] [output_image] [width] [height] [steps] [seed] [use_cfg] [cfg_scale]\n");
|
||||
MNN_PRINT("\n");
|
||||
MNN_PRINT("参数说明:\n");
|
||||
MNN_PRINT(" resource_path : 模型资源路径（包含llm/connector/projector/transformer/vae等模型）\n");
|
||||
MNN_PRINT(" mode : 'text2img' 文生图模式, 'img2img' 图像编辑模式\n");
|
||||
MNN_PRINT(" prompt : 文本描述（支持复杂语义，由Qwen3-0.6B处理）\n");
|
||||
MNN_PRINT(" input_image : 输入图像路径（img2img模式必需，text2img模式忽略，默认: \"\"）\n");
|
||||
MNN_PRINT(" output_image : 输出图像路径（默认: sana_out.jpg）\n");
|
||||
MNN_PRINT(" width : 输出图像宽度（默认: 512，必须是32的倍数）\n");
|
||||
MNN_PRINT(" height : 输出图像高度（默认: 512，必须是32的倍数）\n");
|
||||
MNN_PRINT(" steps : 推理步数（默认: 5，通过蒸馏可用较少步数获得高质量）\n");
|
||||
MNN_PRINT(" seed : 随机种子（默认: 42）\n");
|
||||
MNN_PRINT(" use_cfg : 是否使用CFG引导, 0或1（默认: 0）\n");
|
||||
MNN_PRINT(" cfg_scale : CFG引导强度（默认: 4.5，仅use_cfg=1时生效）\n");
|
||||
MNN_PRINT("\n");
|
||||
MNN_PRINT("示例:\n");
|
||||
MNN_PRINT(" 文生图(512x512): ./sana_diffusion_demo models text2img \"一只可爱的猫咪\" \"\" output.jpg 512 512 5 42 1 4.5\n");
|
||||
MNN_PRINT(" 文生图(1024x1024): ./sana_diffusion_demo models text2img \"一只可爱的猫咪\" \"\" output.jpg 1024 1024 5 42 1 4.5\n");
|
||||
MNN_PRINT(" 图像编辑: ./sana_diffusion_demo models img2img \"添加彩虹\" input.jpg output.jpg 512 512 5 42 0 4.5\n");
|
||||
MNN_PRINT("=====================================================================================================================\n");
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::string resource_path = argv[1];
|
||||
std::string mode = argv[2];
|
||||
std::string prompt = argv[3];
|
||||
std::string image_path = (argc > 4) ? argv[4] : "";
|
||||
std::string output_name = (argc > 5) ? argv[5] : "sana_out.jpg";
|
||||
int width = (argc > 6) ? atoi(argv[6]) : 512;
|
||||
int height = (argc > 7) ? atoi(argv[7]) : 512;
|
||||
int steps = (argc > 8) ? atoi(argv[8]) : 5;
|
||||
int seed = (argc > 9) ? atoi(argv[9]) : 42;
|
||||
bool use_cfg = (argc > 10) ? (atoi(argv[10]) != 0) : false;
|
||||
float cfg_scale = (argc > 11) ? atof(argv[11]) : 4.5f;
|
||||
|
||||
int memory_mode = 2; // standard，0:卸载
|
||||
int backend_type = MNN_FORWARD_CPU;
|
||||
|
||||
// 验证mode参数
|
||||
if (mode != "text2img" && mode != "img2img") {
|
||||
MNN_ERROR("Error: mode must be 'text2img' or 'img2img'\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// 验证img2img模式需要输入图像
|
||||
if (mode == "img2img" && image_path.empty()) {
|
||||
MNN_ERROR("Error: img2img mode requires input image path\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
MNN_PRINT("\n========== 配置信息 ==========\n");
|
||||
MNN_PRINT("模式: %s\n", mode.c_str());
|
||||
MNN_PRINT("提示词: %s\n", prompt.c_str());
|
||||
if (mode == "img2img") {
|
||||
MNN_PRINT("输入图像: %s\n", image_path.c_str());
|
||||
}
|
||||
MNN_PRINT("输出图像: %s\n", output_name.c_str());
|
||||
MNN_PRINT("输出分辨率: %dx%d\n", width, height);
|
||||
MNN_PRINT("推理步数: %d \n", steps);
|
||||
MNN_PRINT("随机种子: %d\n", seed);
|
||||
MNN_PRINT("使用CFG: %s\n", use_cfg ? "是" : "否");
|
||||
if (use_cfg) {
|
||||
MNN_PRINT("CFG强度: %.2f\n", cfg_scale);
|
||||
}
|
||||
MNN_PRINT("==============================\n\n");
|
||||
|
||||
// ========== 步骤1: 初始化Qwen3-0.6B LLM ==========
|
||||
MNN_PRINT("[1/4] 初始化Qwen3-0.6B LLM（文本编码器）...\n");
|
||||
std::string llm_path = resource_path + "/llm/";
|
||||
SanaLlm sana_llm(llm_path);
|
||||
|
||||
// ========== 步骤2: 初始化Diffusion模型 ==========
|
||||
MNN_PRINT("[2/4] 初始化Diffusion模型（Connector + Projector + DiT + VAE）...\n");
|
||||
std::unique_ptr<Diffusion> diffusion(Diffusion::createDiffusion(
|
||||
resource_path,
|
||||
SANA_DIFFUSION,
|
||||
(MNNForwardType)backend_type,
|
||||
memory_mode
|
||||
));
|
||||
diffusion->load();
|
||||
|
||||
// ========== 步骤3: LLM处理文本 ==========
|
||||
MNN_PRINT("[3/4] LLM处理文本prompt...\n");
|
||||
VARP llm_out;
|
||||
|
||||
if (use_cfg) {
|
||||
// CFG模式：同时生成正负样本特征（batch_size=2）
|
||||
MNN_PRINT(" CFG模式：生成正负样本特征\n");
|
||||
llm_out = sana_llm.process(prompt, true, "");
|
||||
} else {
|
||||
// 非CFG模式：只生成正样本特征（batch_size=1）
|
||||
MNN_PRINT(" 非CFG模式：生成单一特征\n");
|
||||
llm_out = sana_llm.process(prompt, false);
|
||||
}
|
||||
|
||||
if (llm_out.get() == nullptr) {
|
||||
MNN_ERROR("LLM处理失败\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// ========== 步骤4: Diffusion生成图像 ==========
|
||||
MNN_PRINT("[4/4] Diffusion生成图像...\n");
|
||||
auto progress = [](int p) {
|
||||
std::cout << " 生成进度: " << p << "%\r" << std::flush;
|
||||
};
|
||||
|
||||
bool success = diffusion->run(
|
||||
llm_out, // LLM特征（来自Qwen3-0.6B）
|
||||
mode, // 模式：text2img或img2img
|
||||
image_path, // 输入图像（img2img模式）
|
||||
output_name, // 输出路径
|
||||
width, // 输出宽度
|
||||
height, // 输出高度
|
||||
steps, // 推理步数（蒸馏加速）
|
||||
seed, // 随机种子
|
||||
use_cfg, // 是否使用CFG
|
||||
cfg_scale, // CFG强度
|
||||
progress // 进度回调
|
||||
);
|
||||
|
||||
if (success) {
|
||||
MNN_PRINT("\n\n========== 生成完成 ==========\n");
|
||||
MNN_PRINT("✓ 图像已保存至: %s\n", output_name.c_str());
|
||||
MNN_PRINT("==============================\n");
|
||||
} else {
|
||||
MNN_ERROR("\n生成失败\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user