Merge pull request #4 from blindcrone/webend

list.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+export AWS_SECRET_ACCESS_KEY=$(cat .env.local | grep AWS_SECRET | cut -d "\"" -f 2)
+export AWS_ACCESS_KEY=$(cat .env.local | grep AWS_ACCESS | cut -d "\"" -f 2)
+aws s3 ls models.webaverse.com | grep glb | cut -d " " -f 6 | cut -d "." -f 1
+
+

main.py

@@ -157,4 +157,4 @@ if __name__ == '__main__':
             trainer.test(test_loader)
 
             if opt.save_mesh:
-                trainer.save_mesh(resolution=256)
+                trainer.save_mesh(resolution=256)

nerf/renderer.py

@@ -296,8 +296,11 @@ class NeRFRenderer(nn.Module):
                 fp.write(f'Ns 0.000000 \n')
                 fp.write(f'map_Kd {name}albedo.png \n')
 
+            os.system(f"blender -b -P obj2glb.py")
+
         _export(v, f)
 
+
     def run(self, rays_o, rays_d, num_steps=128, upsample_steps=128, light_d=None, ambient_ratio=1.0, shading='albedo', bg_color=None, perturb=False, **kwargs):
         # rays_o, rays_d: [B, N, 3], assumes B == 1
         # bg_color: [BN, 3] in range [0, 1]
@@ -642,4 +645,4 @@ class NeRFRenderer(nn.Module):
         else:
             results = _run(rays_o, rays_d, **kwargs)
 
-        return results
+        return results

obj2glb.py

@@ -0,0 +1,11 @@
+# using blender, open an input.obj file, and export it as output.glb
+import bpy
+# import input.obj file
+bpy.ops.import_scene.obj(filepath='trial/mesh/mesh.obj')
+
+# select all
+bpy.ops.object.select_all(action='SELECT')
+
+# export output.glb file
+bpy.ops.export_scene.gltf(filepath="model.glb")
+print('done')

server.cpp

@@ -0,0 +1,166 @@
+#include <cstdio>
+#include <iostream>
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <thread>
+#include <array>
+#include <queue>
+#include <mutex>
+#include <vector>
+#include <filesystem>
+#include <crow.h>
+#include <ranges>
+#include "threadpool.h"
+#include <sstream>
+#include <chrono>
+#include <unordered_map>
+#define CROW_MAIN
+
+using std::string;
+using std::mutex;
+using std::lock_guard;
+using std::make_shared;
+using std::queue;
+using std::vector;
+namespace fs = std::filesystem;
+namespace rv = std::ranges::views;
+
+constexpr string to_st(const auto& i){ std::stringstream ss; ss << i; return ss.str(); }
+template<typename T>
+constexpr T st_to(const string& s){
+	T t(0);
+	std::stringstream ss(s);
+	ss >> t;
+	return t;
+}
+
+static inline string uid(const std::string& s){
+	std::stringstream ss;
+	std::hash<string> h;
+	const auto t0 = std::chrono::system_clock::now();
+	ss << s << '|' << std::chrono::duration_cast<std::chrono::nanoseconds>(t0.time_since_epoch()).count();
+	return to_st(h(ss.str()));
+}
+
+static inline string exec(const char* cmd) {
+	std::array<char, 128> buffer;
+	std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd, "r"), pclose);
+	string result;
+	if (!pipe)
+		return "Command failed";
+	else {
+		while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr)
+			result += buffer.data();
+	}
+	return result;
+}
+
+constexpr auto reify(std::ranges::range auto&& r){ return vector(r.begin(), r.end()); }
+
+constexpr string strip(const string& s){ return s; }
+
+constexpr string strip(string s, char ch){
+	s.erase(std::remove_if(s.begin(), s.end(), [=](char c){ return c == ch; }), s.end());
+	return s;
+}
+
+constexpr string strip(string s, char ch, auto... chs){
+	return strip(strip(s, ch), chs...);
+}
+
+constexpr vector<string> splitOn(const string& s, const string& delim){
+	vector<string> ret;
+	long long int start = 0;
+	for(size_t dp = s.find_first_of(delim, 0); start >= 0
+	   ; start = (dp == std::string::npos) ? -1 : dp + 1, dp = s.find_first_of(delim, start))
+		if(auto n = s.substr(start, dp - start); !n.empty())
+			ret.emplace_back(n);
+	return ret;
+}
+
+template <typename T>
+constexpr auto q_to_v(queue<T> qcopy){
+	vector<T> v;
+	v.reserve(qcopy.size());
+	while(!qcopy.empty())
+		v.push_back(qcopy.front()), qcopy.pop();
+	return v;
+}
+
+int main(){
+	crow::SimpleApp app;
+	typedef std::array<string, 2> guy;
+	auto commissions = make_shared<queue<guy>>();
+	auto queue_mutex = make_shared<mutex>()
+	   , train_mutex = make_shared<mutex>();
+	auto pool = make_shared<threadpool<>>(avail_threads() / 2);
+
+	auto run = [=](const string& cmd){
+		CROW_LOG_INFO << "running \'" << cmd;
+		return exec(cmd.c_str());
+	};
+
+	auto poppe = [=](){
+		lock_guard<mutex> qlock(*queue_mutex);
+		commissions->pop();
+		CROW_LOG_INFO << commissions->size() << " left in queue";
+	};
+
+	auto training_loop = [=](){
+		lock_guard<mutex> lock(*train_mutex);
+		while(!commissions->empty()){
+			auto& [id, prompt] = commissions->front();
+			CROW_LOG_INFO << "Launched training for prompt: " + prompt;
+			run(string("sh train.sh \"") + prompt + "\"");
+			CROW_LOG_INFO << run(string("sh upload.sh ") + id);
+			CROW_LOG_INFO << "Finished training for prompt: " + prompt;
+			poppe();
+		}
+	};
+
+	auto enqueue = [=](const guy& thing){
+		lock_guard<mutex> lock(*queue_mutex);
+		commissions->push(thing);
+		auto& [name, prompt] = thing;
+		CROW_LOG_INFO << name << " queued with prompt: " << prompt;
+	};
+
+	CROW_ROUTE(app, "/create")([=](const crow::request& req){
+		crow::json::wvalue ret;
+		if(auto prompt = req.url_params.get("prompt"); prompt == nullptr){
+			CROW_LOG_INFO << "No prompt specified";
+			ret["error"] = "No prompt given";
+		} else {
+			CROW_LOG_INFO << prompt << " commissioned";
+			auto id = uid(prompt);
+			enqueue({to_st(id), strip(prompt, '\'', '\"')});
+			pool->enqueue(training_loop);
+			CROW_LOG_INFO << "Launched training loop";
+			ret["id"] = id;
+		}
+		return ret;
+	});
+
+	CROW_ROUTE(app, "/list")([=](){
+		auto l = splitOn(run("sh list.sh"), "\n");
+		crow::json::wvalue ret;
+		for(int k = 0; auto& [i, p] : q_to_v(*commissions))
+			ret["pending"][k++] = {{ "id", i }, { "prompt", p}};
+		ret["finished"] = l;
+		return ret;
+	});
+
+	CROW_ROUTE(app, "/get/<string>")([=](const crow::request& req, crow::response& res, const string n){
+		if(auto l = reify( splitOn(run("sh list.sh"), "\n") | rv::filter([=](const string& s){ return s == n; })); !l.empty())
+			res.redirect(string("https://s3.us-west-2.amazonaws.com/models.webaverse.com/") + n + ".glb");
+		else if(auto q = reify( q_to_v(*commissions)
+							   | rv::filter([=](const guy& i){ return i[0] == n; })); !q.empty())
+			res.code = 209;
+		else
+			res.code = 404;
+		res.end();
+	});
+
+	app.port(80).run();
+}

threadpool.h

@@ -0,0 +1,145 @@
+#pragma once
+#include <atomic>
+#include <condition_variable>
+#include <functional>
+#include <future>
+#include <mutex>
+#include <queue>
+#include <utility>
+#include <vector>
+#include <concepts>
+#include <thread>
+static inline size_t avail_threads(){
+	return std::thread::hardware_concurrency();
+}
+
+template <typename F, typename... Args>
+using irt = std::invoke_result_t<F, Args...>;
+
+template <typename F, typename... Args>
+static inline auto taskify(const F& funk, Args... args){
+	return std::bind(funk, std::forward<Args...>(args)...);
+}
+
+static inline auto promitask(const std::invocable<> auto& funk){
+	typedef irt<decltype(funk)> ret;
+	auto pr = std::make_shared<std::promise<ret>>();
+	if constexpr (std::same_as<irt<decltype(funk)>, void>)
+		return make_pair([=](){ funk(), pr->set_value(); }, pr);
+	else
+		return make_pair([=](){ pr->set_value(funk()); }, pr);
+}
+
+static inline auto await_many(const std::ranges::range auto& fu){
+	std::for_each(fu.begin(), fu.end(), [](auto&& f){ f.wait(); });
+}
+
+typedef std::function<void()> tasque;
+
+template <typename Thrd = std::jthread>
+class threadpool {
+	/// If true the queue thread should exit
+	std::atomic<bool> done;
+
+	/// The thread object associated with this queue
+	std::vector<Thrd> queue_threads;
+	/// A queue of functions that will be executed on the queue thread
+	std::queue<tasque> work_queue;
+
+	/// The mutex used in the condition variable
+	std::mutex queue_mutex;
+
+	/// The condition variable that waits for a new function to be inserted in the
+	/// queue
+	std::condition_variable cond;
+
+	/// This funciton executes on the queue_thread
+	void queue_runner() {
+		while (!done) {
+			tasque func;
+			{
+				std::unique_lock<std::mutex> lock(queue_mutex);
+				cond.wait( lock
+						 , [this]() { return work_queue.empty() == false || done; });
+
+				if (!done){
+					swap(func, work_queue.front());
+					work_queue.pop();
+				}
+			}
+			if (func) func();
+		}
+	}
+
+	void qup(const std::invocable<> auto& f){
+		std::lock_guard<std::mutex> lock(queue_mutex);
+		work_queue.push(f);
+		cond.notify_one();
+	}
+
+public:
+	template <typename F, typename... Args>
+	void enqueue(const F& func, Args... args) requires std::invocable<F, Args...> {
+		qup(taskify(func, args...));
+	}
+
+	template <typename F, typename... Args>
+	auto inquire(const F& func, Args... args) requires std::invocable<F, Args...> {
+		auto [t, pr] = promitask(taskify(func, args...));
+		auto fut = pr->get_future();
+		enqueue(t);
+		return fut;
+	}
+
+	void clear() {
+		{
+			std::lock_guard<std::mutex> lock(queue_mutex);
+			while(!work_queue.empty())
+				work_queue.pop();
+		}
+		sync();
+	}
+
+	void sync(){
+		std::atomic<size_t> n(0);
+		const size_t m = queue_threads.size();
+		auto present = [&](){ ++n; size_t l = n.load(); while(l < m) l = n.load(); };
+		std::vector<std::future<void>> fu;
+		std::ranges::generate_n(std::back_inserter(fu), m, [=, this](){ return inquire(present); });
+		await_many(fu);
+	}
+
+	threadpool(size_t n, size_t res) : done(false)
+									 , queue_threads(n ? std::clamp(n, size_t(1), avail_threads() - res)
+													   : std::max(size_t(1), avail_threads() - res)) {
+		for(auto& i:queue_threads){
+			Thrd tmp(&threadpool::queue_runner, this);
+			std::swap(i, tmp);
+		}
+	}
+	threadpool(size_t n) : threadpool(n, 0) {}
+	threadpool() : threadpool(0, 1) {}
+
+	~threadpool() {
+		sync();
+		done.store(true);
+		cond.notify_all();
+	}
+
+	threadpool(const threadpool& other) : work_queue(other.work_queue), done(false) {
+		for(auto& i:queue_threads){
+			Thrd tmp(&threadpool::queue_runner, this);
+			std::swap(i, tmp);
+		}
+	}
+
+	threadpool& operator=(const threadpool& other){
+		clear();
+		work_queue = other.work_queue;
+		return *this;
+	}
+	size_t size() const { return queue_threads.size(); }
+	threadpool& operator=(threadpool&& other) = default;
+	threadpool(threadpool&& other) = default;
+};
+

train.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+rm -rf trial/checkpoints/*
+python main.py --save_mesh --text "$1" --workspace trial -O

upload.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+export AWS_SECRET_ACCESS_KEY=$(cat .env.local | grep AWS_SECRET | cut -d "\"" -f 2)
+export AWS_ACCESS_KEY=$(cat .env.local | grep AWS_ACCESS | cut -d "\"" -f 2)
+aws s3 cp model.glb 's3://models.webaverse.com/'$1'.glb'
+rm model.glb
+