imagen-google/main.py at main · YkywzCode/imagen-google · GitHub

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
import os.path
import sys

from PySide6.QtCore import QUrl
from PySide6.QtGui import QDesktopServices, QIcon
from PySide6.QtWidgets import QApplication, QMainWindow, QPushButton, QVBoxLayout, QWidget, QTextEdit, QSpinBox, QLabel, QComboBox, QMessageBox

from thread import GenerateThread

class ImageGeneratorApp(QMainWindow):
    def __init__(self, base_dir):
        super().__init__()

        self.output = os.path.join(base_dir, 'output/')
        os.makedirs(os.path.dirname(self.output), exist_ok=True)

        self.setWindowTitle("Image Generator with Google Imagen")
        self.resize(450, 450)

        self.text_label = QLabel('Prompt')
        self.text_prompt = QTextEdit()
        self.text_prompt.setAcceptRichText(False)
        self.text_prompt.setPlaceholderText("Enter your image prompt here...")

        self.total_label = QLabel('Total Image')
        self.total_input = QSpinBox()
        self.total_input.setValue(1)
        self.total_input.setRange(1, 100)

        self.aspect_ratio_label = QLabel('Aspect Ratio')
        self.aspect_ratio_input = QComboBox()
        self.aspect_ratio_input.addItems([
            '9:16',
            '16:9'
        ])

        self.generate_button = QPushButton("Generate Image", self)
        self.generate_button.setCheckable(True)
        self.generate_button.clicked.connect(self.generate_image)

        self.output_button = QPushButton('Output')
        self.output_button.clicked.connect(lambda: QDesktopServices.openUrl(QUrl.fromLocalFile(self.output)))

        layout = QVBoxLayout()
        layout.addWidget(self.text_label)
        layout.addWidget(self.text_prompt)
        layout.addWidget(self.total_label)
        layout.addWidget(self.total_input)
        layout.addWidget(self.aspect_ratio_label)
        layout.addWidget(self.aspect_ratio_input)
        layout.addWidget(self.generate_button)
        layout.addWidget(self.output_button)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def generate_image(self):
        text_prompt = self.text_prompt.toPlainText()
        total_generate = self.total_input.value()
        aspect_ratio = self.aspect_ratio_input.currentText()
        if not text_prompt:
            print("Please enter a prompt.")
            return

        def started():
            self.generate_button.setDisabled(True)
            self.generate_button.setText('Loading')
        def finished():
            self.generate_button.setChecked(False)
            self.generate_button.setDisabled(False)
            self.generate_button.setText('Generate')
        def result(val):
            QMessageBox.information(self, 'Info', 'finish generate\n'
                                                  f'succes: {val["success"]}\n'
                                                  f'error: {val["error"]}')

        if self.generate_button.isChecked():
            self.gen_thread = GenerateThread(self, {
                'prompt': text_prompt,
                'total': total_generate,
                'ar': aspect_ratio,
                'output': self.output
            })
            self.gen_thread.started.connect(started)
            self.gen_thread.finished.connect(finished)
            self.gen_thread.result.connect(result)
            self.gen_thread.start()
        else:
            self.gen_thread.terminate()


if __name__ == "__main__":
    app = QApplication(sys.argv)

    if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
        base_dir = os.path.dirname(sys.executable)
        application_path = sys._MEIPASS
    elif __file__:
        base_dir = os.path.abspath(".")
        application_path = os.path.abspath(".")

    icon = QIcon(os.path.join(application_path, 'app.ico'))
    app.setWindowIcon(icon)

    window = ImageGeneratorApp(base_dir)
    window.show()
    sys.exit(app.exec())