File size: 4,669 Bytes
b40fda4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 discord
import logging
import os
from io import BytesIO
from gradio_client import Client

# λ‘œκΉ… μ„€μ •
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# μΈν…νŠΈ μ„€μ •
intents = discord.Intents.default()
intents.message_content = True

# API ν΄λΌμ΄μ–ΈνŠΈ μ„€μ •
api_client = Client("http://211.233.58.202:7960/")

# λ””μŠ€μ½”λ“œ 봇 클래슀
class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False

    async def on_ready(self):
        logging.info(f'{self.user}둜 λ‘œκ·ΈμΈλ˜μ—ˆμŠ΅λ‹ˆλ‹€!')
        # μƒν˜Έ λŒ€ν™” κΈ°λŠ₯ μ•Œλ¦Ό
        channel = self.get_channel(int(os.getenv("DISCORD_CHANNEL_ID", "123456789012345678")))
        await channel.send("μ €λŠ” 이미지 생성을 μˆ˜ν–‰ν•  수 있으며, μƒμ„±λœ 이미지에 λŒ€ν•œ μ„€λͺ…을 ν•œκΈ€λ‘œ μ œκ³΅ν•˜κ³  μƒν˜Έ λŒ€ν™”λ₯Ό ν•  수 μžˆμŠ΅λ‹ˆλ‹€. '!image <ν”„λ‘¬ν”„νŠΈ>'λ₯Ό μ‚¬μš©ν•˜μ—¬ 이미지λ₯Ό μš”μ²­ν•˜μ„Έμš”.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if message.content.startswith('!image '):
            if self.is_processing:
                await message.channel.send("이미지 생성이 이미 μ§„ν–‰ μ€‘μž…λ‹ˆλ‹€. μž μ‹œλ§Œ κΈ°λ‹€λ € μ£Όμ„Έμš”.")
                return
            self.is_processing = True
            try:
                prompt = message.content[len('!image '):]
                image_path, used_seed, translated_prompt = await self.generate_image(prompt)
                user_id = message.author.id
                await message.channel.send(
                    f"<@{user_id}> λ‹˜μ΄ μš”μ²­ν•˜μ‹  μ΄λ―Έμ§€μž…λ‹ˆλ‹€.\n"
                    f"μ‚¬μš©λœ μ‹œλ“œ: {used_seed}\n"
                    f"λ²ˆμ—­λœ ν”„λ‘¬ν”„νŠΈ: {translated_prompt}",
                    file=discord.File(image_path)
                )
                # 이미지 생성 ν›„ μ„€λͺ… 제곡 및 λŒ€ν™”
                await initiate_conversation(prompt, message)
            except Exception as e:
                logging.error(f'이미지 생성 쀑 였λ₯˜ λ°œμƒ: {e}')
                await message.channel.send("μ‚¬μš© μ˜ˆμ‹œ: !image 고양이가 'HAPPY WORLD'라고 μ“°μ—¬μ§„ ν‘œμ§€νŒμ„ λ“€κ³ μžˆλ‹€. ")
            finally:
                self.is_processing = False    

    async def generate_image(self, prompt):
        if not prompt:
            raise ValueError("Prompt is empty or None")
    
        logging.debug(f"Sending request to API with prompt: {prompt}")

        try:
            result = api_client.predict(
                prompt=prompt,
                seed=123,
                randomize_seed=False,
                width=1024,
                height=576,
                guidance_scale=5,
                num_inference_steps=28,
                api_name="/infer_t2i"
            )
            logging.debug(f"API response received: {result}")
        
            if isinstance(result, tuple) and len(result) == 3:
                image_path, used_seed, translated_prompt = result
                logging.info(f"Image generated at: {image_path}")
                logging.info(f"Used seed: {used_seed}")
                logging.info(f"Translated prompt: {translated_prompt}")
            
                return image_path, used_seed, translated_prompt
            else:
                raise ValueError("Unexpected API response format")
        except gradio_client.exceptions.AppError as e:
            logging.error(f"Gradio App Error: {e}")
            raise RuntimeError(f"API μ„œλ²„ 였λ₯˜: {str(e)}. μ„œλ²„ κ΄€λ¦¬μžμ—κ²Œ λ¬Έμ˜ν•˜μ„Έμš”.")
        except Exception as e:
            logging.error(f'API μš”μ²­ 쀑 였λ₯˜ λ°œμƒ: {e}', exc_info=True)
            raise RuntimeError(f"API μš”μ²­ 쀑 였λ₯˜ λ°œμƒ: {str(e)}")
            

async def initiate_conversation(prompt, message):
    logging.debug(f'λŒ€ν™” μ‹œμž‘ 쀑: {prompt}')
    # 이미지 μ„€λͺ…을 ν•œκΈ€λ‘œ 생성
    description = "μƒμ„±λœ μ΄λ―Έμ§€μž…λ‹ˆλ‹€."
    logging.debug(f'이미지 μ„€λͺ…: {description}')

    await message.channel.send(f"이미지 μ„€λͺ…: {description}")
    await continue_conversation(prompt, message)

async def continue_conversation(prompt, message):
    # λŒ€ν™” 지속 κΈ°λŠ₯
    logging.debug(f'λŒ€ν™” 지속 쀑: {prompt}')
    # 계속 λŒ€ν™” λ‚΄μš©μ„ λ°›μ•„ μƒν˜Έμž‘μš©ν•˜λ„λ‘ κ΅¬ν˜„

# λ””μŠ€μ½”λ“œ 토큰 및 봇 μ‹€ν–‰
if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    discord_client = MyClient(intents=intents)
    discord_client.run(discord_token)