-
Notifications
You must be signed in to change notification settings - Fork 736
Expand file tree
/
Copy pathmain.cpp
More file actions
158 lines (134 loc) · 4.91 KB
/
Copy pathmain.cpp
File metadata and controls
158 lines (134 loc) · 4.91 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
#include <cstdlib>
#include <iostream>
#include <mutex>
#include <string>
#include <thread>
#include <vector>
#include "httplib.h"
#include "async_jobs.h"
#include "common/common.h"
#include "common/resource_owners.hpp"
#include "routes.h"
#include "runtime.h"
#ifdef HAVE_INDEX_HTML
#include "frontend/dist/gen_index_html.h"
#endif
static void print_usage(const char* argv0, const std::vector<ArgOptions>& options_list) {
std::cout << version_string() << "\n";
std::cout << "Usage: " << argv0 << " [options]\n\n";
std::cout << "Svr Options:\n";
options_list[0].print();
std::cout << "\nContext Options:\n";
options_list[1].print();
std::cout << "\nDefault Generation Options:\n";
options_list[2].print();
}
static void parse_args(int argc,
const char** argv,
SDSvrParams& svr_params,
SDContextParams& ctx_params,
SDGenerationParams& default_gen_params) {
std::vector<ArgOptions> options_vec = {
svr_params.get_options(),
ctx_params.get_options(),
default_gen_params.get_options(),
};
if (!parse_options(argc, argv, options_vec)) {
print_usage(argv[0], options_vec);
exit(svr_params.normal_exit ? 0 : 1);
}
const bool random_seed_requested = default_gen_params.seed < 0;
if (!svr_params.resolve_and_validate() ||
!ctx_params.resolve_and_validate(IMG_GEN) ||
!default_gen_params.resolve_and_validate(IMG_GEN,
ctx_params.lora_model_dir,
ctx_params.hires_upscalers_dir)) {
print_usage(argv[0], options_vec);
exit(1);
}
if (random_seed_requested) {
default_gen_params.seed = -1;
}
}
void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
SDSvrParams* svr_params = (SDSvrParams*)data;
log_print(level, log, svr_params->verbose, svr_params->color);
}
int main(int argc, const char** argv) {
if (argc > 1 && std::string(argv[1]) == "--version") {
std::cout << version_string() << "\n";
return EXIT_SUCCESS;
}
SDSvrParams svr_params;
SDContextParams ctx_params;
SDGenerationParams default_gen_params;
parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
sd_set_log_callback(sd_log_cb, (void*)&svr_params);
log_verbose = svr_params.verbose;
log_color = svr_params.color;
LOG_DEBUG("version: %s", version_string().c_str());
LOG_DEBUG("%s", sd_get_system_info());
LOG_DEBUG("%s", svr_params.to_string().c_str());
LOG_DEBUG("%s", ctx_params.to_string().c_str());
LOG_DEBUG("%s", default_gen_params.to_string().c_str());
sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false);
SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
if (sd_ctx == nullptr) {
LOG_ERROR("new_sd_ctx_t failed");
return 1;
}
std::mutex sd_ctx_mutex;
std::vector<LoraEntry> lora_cache;
std::mutex lora_mutex;
std::vector<UpscalerEntry> upscaler_cache;
std::mutex upscaler_mutex;
AsyncJobManager async_job_manager;
ServerRuntime runtime = {
sd_ctx.get(),
&sd_ctx_mutex,
&svr_params,
&ctx_params,
&default_gen_params,
&lora_cache,
&lora_mutex,
&upscaler_cache,
&upscaler_mutex,
&async_job_manager,
};
std::thread async_worker(async_job_worker, std::ref(runtime));
httplib::Server svr;
svr.set_pre_routing_handler([](const httplib::Request& req, httplib::Response& res) {
std::string origin = req.get_header_value("Origin");
if (origin.empty()) {
origin = "*";
}
res.set_header("Access-Control-Allow-Origin", origin);
res.set_header("Access-Control-Allow-Credentials", "true");
res.set_header("Access-Control-Allow-Methods", "*");
res.set_header("Access-Control-Allow-Headers", "*");
if (req.method == "OPTIONS") {
res.status = 204;
return httplib::Server::HandlerResponse::Handled;
}
return httplib::Server::HandlerResponse::Unhandled;
});
std::string index_html;
#ifdef HAVE_INDEX_HTML
index_html.assign(reinterpret_cast<const char*>(index_html_bytes), index_html_size);
#else
index_html = "Stable Diffusion Server is running";
#endif
register_index_endpoints(svr, svr_params, index_html);
register_openai_api_endpoints(svr, runtime);
register_sdapi_endpoints(svr, runtime);
register_sdcpp_api_endpoints(svr, runtime);
LOG_INFO("listening on: http://%s:%d\n", svr_params.listen_ip.c_str(), svr_params.listen_port);
svr.listen(svr_params.listen_ip, svr_params.listen_port);
{
std::lock_guard<std::mutex> lock(async_job_manager.mutex);
async_job_manager.stop = true;
}
async_job_manager.cv.notify_all();
async_worker.join();
return 0;
}