app.py

import requests
import discord
import logging
import os
from transformers import pipeline

# 로깅 설정
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

# 번역 파이프라인 설정
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# 고정된 네거티브 프롬프트
negative_prompt = "blur, low quality, bad composition, ugly, disfigured, weird colors, low quality, jpeg artifacts, lowres, grainy, deformed structures, blurry, opaque, low contrast, distorted details, details are low"

# 인퍼런스 API를 사용하기 위한 함수
def generate_image(prompt, negative_prompt):
    headers = {
        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
    }
    data = {
        "inputs": prompt + " " + negative_prompt  # 프롬프트와 네거티브 프롬프트를 하나의 문자열로 결합
    }
    api_url = "https://api-inference.huggingface.co/models/fluently/Fluently-XL-Final"
    response = requests.post(api_url, headers=headers, json=data)
    if response.status_code == 200:
        image_url = response.json()[0]['url']  # 응답에서 이미지 URL 추출
        return image_url
    else:
        logging.error(f"API 요청에 실패했습니다: {response.text}")
        return None

# 프롬프트 번역 함수
def translate_prompt(prompt):
    logging.debug(f'프롬프트 번역 중: {prompt}')
    translation = translator(prompt, max_length=512)
    translated_text = translation[0]['translation_text']
    logging.debug(f'번역된 텍스트: {translated_text}')
    return translated_text

# 디스코드 봇 클래스
class MyClient(discord.Client):
    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')

    async def on_message(self, message):
        if message.author == self.user:
            return
        if message.content.startswith('!image '):
            self.is_processing = True
            try:
                prompt = message.content[len('!image '):]
                prompt_en = translate_prompt(prompt)
                logging.debug(f'번역된 프롬프트: {prompt_en}')
                logging.debug(f'고정된 네거티브 프롬프트: {negative_prompt}')

                image_url = generate_image(prompt_en, negative_prompt)
                user_id = message.author.id
                if image_url:
                    await message.channel.send(
                        f"<@{user_id}> 님이 요청하신 이미지입니다: {image_url}"
                    )
                else:
                    await message.channel.send(f"<@{user_id}> 이미지 생성에 실패하였습니다.")
            finally:
                self.is_processing = False
        else:
            # "!image" 명령어가 아닐 경우 안내 메시지 전송
            await message.channel.send('올바른 명령어를 입력해 주세요. 예) "!image 귀여운 고양이가 잠을 자고있다." 등으로 입력하시면 이미지가 생성됩니다.')

# 디스코드 토큰 및 봇 실행
if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    discord_client = MyClient(intents=intents)
    discord_client.run(discord_token)