app.py

import requests
import discord
import logging
import os
from transformers import pipeline
import subprocess
import torch
from diffusers import DiffusionPipeline
import io
from PIL import Image
import time
from dotenv import load_dotenv

# .env 파일에서 환경 변수 로드
load_dotenv()

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

# 환경 변수에서 지정된 채널 ID 가져오기, 기본값 사용
TARGET_CHANNEL_ID = int(os.getenv('TARGET_CHANNEL_ID', '123456789012345678'))  # 기본값 사용, 실제 채널 ID로 대체 필요

# 고정된 네거티브 프롬프트
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"

# 디바이스 설정
device = "cuda" if torch.cuda.is_available() else "cpu"

# 이미지 생성 파이프라인 설정
pipeline = DiffusionPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float32)
pipeline = pipeline.to(device)

# 프롬프트 번역 함수
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

def generate_image(prompt, negative_prompt):
    combined_prompt = f"{prompt}. {negative_prompt}"
    result = pipeline(combined_prompt)
    image = result.images[0]  # 첫 번째 이미지 선택
    return image

class MyClient(discord.Client):
    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        subprocess.Popen(["python", "web.py"])  # 별도의 Python 스크립트 실행
        logging.info("web.py 서버가 시작되었습니다.")

    async def on_message(self, message):
        if message.author == self.user or message.channel.id != TARGET_CHANNEL_ID:
            return
        if message.content.startswith('!image '):
            self.is_processing = True
            try:
                prompt = message.content[len('!image '):]
                prompt_en = translate_prompt(prompt)
                image = generate_image(prompt_en, negative_prompt)
                user_id = message.author.id
                if image:
                    # 이미지를 Discord에 직접 업로드
                    with io.BytesIO() as image_binary:
                        image.save(image_binary, 'PNG')
                        image_binary.seek(0)
                        await message.channel.send(
                            f"<@{user_id}> 님이 요청하신 이미지입니다:",
                            file=discord.File(fp=image_binary, filename='image.png')
                        )
                else:
                    await message.channel.send(f"<@{user_id}> 이미지 생성에 실패하였습니다.")
            finally:
                self.is_processing = False
        else:
            await message.channel.send('올바른 명령어를 입력해 주세요. 예) "!image 귀여운 고양이가 잠을 자고있다." 등으로 입력하시면 이미지가 생성됩니다.')

# 봇 실행
if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    discord_client = MyClient(intents=intents)
    discord_client.run(discord_token)