gameui/ai
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

通过list建立websocket

Browse Source
This commit is contained in:
桜華
2023-06-20 18:09:22 +08:00
parent 53a5287af3
commit b656065b43
2 changed files with 63 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
models/Image.go
View File

@@ -16,27 +16,27 @@ import (
)
type Image struct {
ID int `json:"id" gorm:"primary_key"`
Name string `json:"name"`
Hash string `json:"hash"`
Path string `json:"path"`
Type string `json:"type"`
Size int `json:"size"`
Width int `json:"width"`
Height int `json:"height"`
Prompt string `json:"prompt"`
Format string `json:"format"`
NegativePrompt string `json:"negative_prompt"`
NumInferenceSteps int `json:"num_inference_steps"` // Number of inference steps (minimum: 1; maximum: 500)
GuidanceScale float32 `json:"guidance_scale"` // Scale for classifier-free guidance (minimum: 1; maximum: 20)
Scheduler string `json:"scheduler"` // (DDIM|K_EULER|DPMSolverMultistep|K_EULER_ANCESTRAL|PNDM|KLMS)
Seed int `json:"seed"` // 随机种子(minimum: 0; maximum: 2147483647)
FromImage int `json:"from_image"` // 来源图片(如果是从图片生成的, 则记录来源图片的ID)
Task string `json:"task"` // 任务编号(uuid)
Status string `json:"status"` // 任务状态(queued|running|finished|failed)
Progress int `json:"progress"` // 任务进度(0-100)
Public bool `json:"public"` // 是否公开
UserID int `json:"user_id"` // 用户ID
ID int `json:"id" gorm:"primary_key"` // ID
Name string `json:"name"` // 名称
Hash string `json:"hash"` // 哈希值
Path string `json:"path"` // 路径
Type string `json:"type"` // 类型
Size int `json:"size"` // 大小
Width int `json:"width"` // 宽度
Height int `json:"height"` // 高度
Format string `json:"format"` // 格式
Prompt string `json:"prompt"` // 提示词
NegativePrompt string `json:"negative_prompt"` // 负向提示
NumInferenceSteps int `json:"num_inference_steps"` // 推理步数(minimum: 1; maximum: 500)
GuidanceScale float32 `json:"guidance_scale"` // 引导比例(minimum: 1; maximum: 20)
Scheduler string `json:"scheduler"` // 调度器(DDIM|K_EULER|DPMSolverMultistep|K_EULER_ANCESTRAL|PNDM|KLMS)
Seed int `json:"seed"` // 随机种子(minimum: 0; maximum: 2147483647)
FromImage int `json:"from_image"` // 来源图片(如果是从图片生成的, 则记录来源图片的ID)
Task string `json:"task"` // 任务编号(uuid)
Status string `json:"status"` // 任务状态(queued|running|finished|failed)
Progress int `json:"progress"` // 任务进度(0-100)
Public bool `json:"public"` // 是否公开
UserID int `json:"user_id"` // 用户ID
CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime"`
UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime"`
}

43
routers/images.go
View File

@@ -19,9 +19,48 @@ import (
"github.com/google/uuid"
"github.com/gorilla/mux"
"github.com/gorilla/websocket"
)
var images_websocket_manager = models.NewWebSocketManager()
func ImagesGet(w http.ResponseWriter, r *http.Request) {
// websocket 推理图像
if r.Header.Get("Upgrade") == "websocket" {
upgrader := websocket.Upgrader{}
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
defer conn.Close()
task := r.URL.Query().Get("task")
if task == "" {
w.WriteHeader(http.StatusBadRequest)
w.Write([]byte("task 参数不能为空"))
return
}
wsid := images_websocket_manager.AddConnection(conn)
defer images_websocket_manager.RemoveConnection(wsid)
for {
_, msg, err := conn.ReadMessage()
if err != nil {
log.Println(err)
return
}
log.Println(string(msg))
if string(msg) == "close" {
break
}
}
return
}
var listview models.ListView
listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
@@ -67,8 +106,10 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
return
}
// 直接创建一组图片
// TODO: 创建任务获得任务编号, 多张图时期望可以流式推理
task := uuid.New().String()
// 直接创建一组图片
var image_list []models.Image
for i := 0; i < template.Number; i++ {
var image models.Image