init v1.0.0

ageer

2024-02-27 a079ef44e53acd9e8df51dbb31cf5aea4f9be5bd

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
package com.xmzs.midjourney.service;
 
import com.xmzs.midjourney.Constants;
import com.xmzs.midjourney.ReturnCode;
import com.xmzs.midjourney.enums.BlendDimensions;
import com.xmzs.midjourney.loadbalancer.DiscordInstance;
import com.xmzs.midjourney.loadbalancer.DiscordLoadBalancer;
import com.xmzs.midjourney.result.Message;
import com.xmzs.midjourney.result.SubmitResultVO;
import com.xmzs.midjourney.support.Task;
import com.xmzs.midjourney.util.MimeTypeUtils;
import eu.maxschuster.dataurl.DataUrl;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
 
import java.util.ArrayList;
import java.util.List;
 
@Slf4j
@Service
@RequiredArgsConstructor
public class TaskServiceImpl implements TaskService {
    private final TaskStoreService taskStoreService;
    private final DiscordLoadBalancer discordLoadBalancer;
 
    @Override
    public SubmitResultVO submitImagine(Task task, List<DataUrl> dataUrls) {
        DiscordInstance instance = this.discordLoadBalancer.chooseInstance();
        if (instance == null) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "无可用的账号实例");
        }
        task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, instance.getInstanceId());
        return instance.submitTask(task, () -> {
            List<String> imageUrls = new ArrayList<>();
            for (DataUrl dataUrl : dataUrls) {
                String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
                Message<String> uploadResult = instance.upload(taskFileName, dataUrl);
                if (uploadResult.getCode() != ReturnCode.SUCCESS) {
                    return Message.of(uploadResult.getCode(), uploadResult.getDescription());
                }
                String finalFileName = uploadResult.getResult();
                Message<String> sendImageResult = instance.sendImageMessage("upload image: " + finalFileName, finalFileName);
                if (sendImageResult.getCode() != ReturnCode.SUCCESS) {
                    return Message.of(sendImageResult.getCode(), sendImageResult.getDescription());
                }
                imageUrls.add(sendImageResult.getResult());
            }
            if (!imageUrls.isEmpty()) {
                task.setPrompt(String.join(" ", imageUrls) + " " + task.getPrompt());
                task.setPromptEn(String.join(" ", imageUrls) + " " + task.getPromptEn());
                task.setDescription("/imagine " + task.getPrompt());
                this.taskStoreService.save(task);
            }
            return instance.imagine(task.getPromptEn(), task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));
        });
    }
 
    @Override
    public SubmitResultVO submitUpscale(Task task, String targetMessageId, String targetMessageHash, int index, int messageFlags) {
        String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
        DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
        if (discordInstance == null || !discordInstance.isAlive()) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
        }
        return discordInstance.submitTask(task, () -> discordInstance.upscale(targetMessageId, index, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
    }
 
    @Override
    public SubmitResultVO submitVariation(Task task, String targetMessageId, String targetMessageHash, int index, int messageFlags) {
        String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
        DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
        if (discordInstance == null || !discordInstance.isAlive()) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
        }
        return discordInstance.submitTask(task, () -> discordInstance.variation(targetMessageId, index, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
    }
 
    @Override
    public SubmitResultVO submitReroll(Task task, String targetMessageId, String targetMessageHash, int messageFlags) {
        String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
        DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
        if (discordInstance == null || !discordInstance.isAlive()) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
        }
        return discordInstance.submitTask(task, () -> discordInstance.reroll(targetMessageId, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
    }
 
    @Override
    public SubmitResultVO submitDescribe(Task task, DataUrl dataUrl) {
        DiscordInstance discordInstance = this.discordLoadBalancer.chooseInstance();
        if (discordInstance == null) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "无可用的账号实例");
        }
        task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, discordInstance.getInstanceId());
        return discordInstance.submitTask(task, () -> {
            String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
            Message<String> uploadResult = discordInstance.upload(taskFileName, dataUrl);
            if (uploadResult.getCode() != ReturnCode.SUCCESS) {
                return Message.of(uploadResult.getCode(), uploadResult.getDescription());
            }
            String finalFileName = uploadResult.getResult();
            return discordInstance.describe(finalFileName, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));
        });
    }
 
    @Override
    public SubmitResultVO submitBlend(Task task, List<DataUrl> dataUrls, BlendDimensions dimensions) {
        DiscordInstance discordInstance = this.discordLoadBalancer.chooseInstance();
        if (discordInstance == null) {
            return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "无可用的账号实例");
        }
        task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, discordInstance.getInstanceId());
        return discordInstance.submitTask(task, () -> {
            List<String> finalFileNames = new ArrayList<>();
            for (DataUrl dataUrl : dataUrls) {
                String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
                Message<String> uploadResult = discordInstance.upload(taskFileName, dataUrl);
                if (uploadResult.getCode() != ReturnCode.SUCCESS) {
                    return Message.of(uploadResult.getCode(), uploadResult.getDescription());
                }
                finalFileNames.add(uploadResult.getResult());
            }
            return discordInstance.blend(finalFileNames, dimensions, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));
        });
    }
 
}