Skip to content

Commit 108c750

Browse files
cmdr2cmdr2
committed
Allow --vae-on-cpu, --vae-tiling, --offload-to-cpu and --diffusion-fa CLI params
1 parent f6beccc commit 108c750

5 files changed

Lines changed: 84 additions & 8 deletions

File tree

‎include/image_generator.h‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
#include "stable-diffusion.h"
1313
#include "task_state.h"
1414

15+
// Forward declaration
16+
struct ServerParams;
17+
1518
struct ImageGenerationParams {
1619
std::string prompt;
1720
std::string negative_prompt;
@@ -45,7 +48,7 @@ class ImageGenerator {
4548
public:
4649
ImageGenerator(std::shared_ptr<TaskStateManager> task_state_manager,
4750
std::shared_ptr<OptionsManager> options_manager, std::shared_ptr<ModelManager> model_manager,
48-
std::shared_ptr<ImageFilters> image_filters);
51+
std::shared_ptr<ImageFilters> image_filters, const ServerParams& server_params);
4952
~ImageGenerator();
5053

5154
// Check if initialized
@@ -106,6 +109,12 @@ class ImageGenerator {
106109
std::string current_lora_model_dir_;
107110
std::string current_embeddings_dir_;
108111
std::string current_controlnet_path_;
112+
113+
// SD context parameters from CLI
114+
bool vae_on_cpu_;
115+
bool vae_tiling_;
116+
bool offload_to_cpu_;
117+
bool diffusion_fa_;
109118
};
110119

111120
#endif // __IMAGE_GENERATOR_H__

‎include/server.h‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,18 @@
1111
#include "options_manager.h"
1212
#include "task_state.h"
1313

14+
struct ServerParams {
15+
int port = 8188;
16+
std::shared_ptr<ModelManager> model_manager;
17+
bool vae_on_cpu = false;
18+
bool vae_tiling = false;
19+
bool offload_to_cpu = false;
20+
bool diffusion_fa = false;
21+
};
22+
1423
class Server {
1524
public:
16-
Server(int port, std::shared_ptr<ModelManager> model_manager);
25+
Server(const ServerParams& params);
1726
~Server();
1827

1928
void run();
@@ -38,6 +47,7 @@ class Server {
3847
// Helper methods
3948
crow::response generateImage(const crow::json::rvalue& json_body, bool is_img2img);
4049

50+
ServerParams params_;
4151
int port_;
4252
crow::SimpleApp app_;
4353
std::shared_ptr<OptionsManager> options_manager_;

‎src/image_generator.cpp‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include "image_utils.h"
1313
#include "logging.h"
1414
#include "model_detection.h"
15+
#include "server.h"
1516

1617
// Global callback data structure
1718
struct CallbackData {
@@ -72,14 +73,19 @@ static void preview_callback(int step, int frame_count, sd_image_t* frames, bool
7273

7374
ImageGenerator::ImageGenerator(std::shared_ptr<TaskStateManager> task_state_manager,
7475
std::shared_ptr<OptionsManager> options_manager,
75-
std::shared_ptr<ModelManager> model_manager, std::shared_ptr<ImageFilters> image_filters)
76+
std::shared_ptr<ModelManager> model_manager, std::shared_ptr<ImageFilters> image_filters,
77+
const ServerParams& server_params)
7678
: sd_ctx_(nullptr),
7779
task_state_manager_(task_state_manager),
7880
options_manager_(options_manager),
7981
model_manager_(model_manager),
8082
image_filters_(image_filters),
8183
initialized_(false),
82-
interrupted_(false) {
84+
interrupted_(false),
85+
vae_on_cpu_(server_params.vae_on_cpu),
86+
vae_tiling_(server_params.vae_tiling),
87+
offload_to_cpu_(server_params.offload_to_cpu),
88+
diffusion_fa_(server_params.diffusion_fa) {
8389
LOG_INFO("ImageGenerator created");
8490
}
8591

@@ -191,6 +197,17 @@ std::vector<std::string> ImageGenerator::generateInternal(const ImageGenerationP
191197
LOG_INFO("Using ControlNet with strength %.2f", params.control_strength);
192198
}
193199

200+
// VAE tiling (if enabled via CLI)
201+
if (vae_tiling_) {
202+
gen_params.vae_tiling_params.enabled = true;
203+
gen_params.vae_tiling_params.tile_size_x = 512; // Default tile size
204+
gen_params.vae_tiling_params.tile_size_y = 512;
205+
gen_params.vae_tiling_params.target_overlap = 0.5f;
206+
gen_params.vae_tiling_params.rel_size_x = 1.0f;
207+
gen_params.vae_tiling_params.rel_size_y = 1.0f;
208+
LOG_INFO("VAE tiling enabled with tile size 512x512");
209+
}
210+
194211
// Generate images
195212
sd_image_t* result = generate_image(sd_ctx_, &gen_params);
196213

@@ -525,6 +542,21 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
525542
// Set RNG type
526543
params.rng_type = CUDA_RNG;
527544

545+
// Apply CLI parameters for SD context
546+
params.keep_vae_on_cpu = vae_on_cpu_;
547+
params.offload_params_to_cpu = offload_to_cpu_;
548+
params.diffusion_flash_attn = diffusion_fa_;
549+
550+
if (vae_on_cpu_) {
551+
LOG_INFO("VAE will be kept on CPU");
552+
}
553+
if (offload_to_cpu_) {
554+
LOG_INFO("Parameters will be offloaded to CPU");
555+
}
556+
if (diffusion_fa_) {
557+
LOG_INFO("Diffusion flash attention enabled");
558+
}
559+
528560
// Create SD context
529561
sd_ctx_ = new_sd_ctx(&params);
530562
if (!sd_ctx_) {

‎src/main.cpp‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,10 @@ struct CommandLineArgs {
8989
std::string embeddings_dir;
9090
std::string controlnet_dir;
9191
std::string text_encoder_dir;
92+
bool vae_on_cpu = false;
93+
bool vae_tiling = false;
94+
bool offload_to_cpu = false;
95+
bool diffusion_fa = false;
9296
};
9397

9498
void print_usage(const char* program_name) {
@@ -109,6 +113,10 @@ void print_usage(const char* program_name) {
109113
std::cerr << " --embeddings-dir <path> Embeddings directory" << std::endl;
110114
std::cerr << " --controlnet-dir <path> ControlNet models directory" << std::endl;
111115
std::cerr << " --text-encoder-dir <path> Text encoder models directory" << std::endl;
116+
std::cerr << " --vae-on-cpu Keep VAE on CPU (default: false)" << std::endl;
117+
std::cerr << " --vae-tiling Enable VAE tiling (default: false)" << std::endl;
118+
std::cerr << " --offload-to-cpu Offload parameters to CPU (default: false)" << std::endl;
119+
std::cerr << " --diffusion-fa Enable diffusion flash attention (default: false)" << std::endl;
112120
}
113121

114122
CommandLineArgs parse_args(int argc, char* argv[]) {
@@ -155,6 +163,14 @@ CommandLineArgs parse_args(int argc, char* argv[]) {
155163
args.controlnet_dir = argv[++i];
156164
} else if (arg == "--text-encoder-dir" && i + 1 < argc) {
157165
args.text_encoder_dir = argv[++i];
166+
} else if (arg == "--vae-on-cpu") {
167+
args.vae_on_cpu = true;
168+
} else if (arg == "--vae-tiling") {
169+
args.vae_tiling = true;
170+
} else if (arg == "--offload-to-cpu") {
171+
args.offload_to_cpu = true;
172+
} else if (arg == "--diffusion-fa") {
173+
args.diffusion_fa = true;
158174
} else if (arg == "--help" || arg == "-h") {
159175
print_usage(argv[0]);
160176
exit(0);
@@ -225,8 +241,17 @@ int main(int argc, char* argv[]) {
225241
}
226242

227243
try {
244+
// Create server parameters
245+
ServerParams server_params;
246+
server_params.port = args.port;
247+
server_params.model_manager = model_manager;
248+
server_params.vae_on_cpu = args.vae_on_cpu;
249+
server_params.vae_tiling = args.vae_tiling;
250+
server_params.offload_to_cpu = args.offload_to_cpu;
251+
server_params.diffusion_fa = args.diffusion_fa;
252+
228253
// Create and start the server
229-
g_server = std::make_unique<Server>(args.port, model_manager);
254+
g_server = std::make_unique<Server>(server_params);
230255
g_server->run();
231256
} catch (const std::exception& e) {
232257
std::cerr << "Server error: " << e.what() << std::endl;

‎src/server.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,8 @@ static int round_to_nearest_multiple_of_64(int dimension) {
107107
return std::round(static_cast<double>(dimension) / 64.0) * 64;
108108
}
109109

110-
Server::Server(int port, std::shared_ptr<ModelManager> model_manager)
111-
: port_(port), model_manager_(model_manager), should_stop_(false) {
110+
Server::Server(const ServerParams& params)
111+
: params_(params), port_(params.port), model_manager_(params.model_manager), should_stop_(false) {
112112
// Set up custom logger to filter out unnecessary requests
113113
static FilteredLogHandler filtered_handler;
114114
crow::logger::setHandler(&filtered_handler);
@@ -121,7 +121,7 @@ Server::Server(int port, std::shared_ptr<ModelManager> model_manager)
121121

122122
// Create ImageGenerator with shared task state manager and model manager
123123
image_generator_ =
124-
std::make_unique<ImageGenerator>(task_state_manager_, options_manager_, model_manager_, image_filters_);
124+
std::make_unique<ImageGenerator>(task_state_manager_, options_manager_, model_manager_, image_filters_, params);
125125

126126
options_manager_->load();
127127

0 commit comments

Comments
 (0)