app.py

import discord
import logging
import os
import uuid
import torch
import subprocess
from huggingface_hub import snapshot_download
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
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

# Hugging Face 모델 다운로드
huggingface_token = os.getenv("HF_TOKEN")
model_path = snapshot_download(
    repo_id="Corcelio/mobius",
    repo_type="model",
    local_dir="mobius",
    token=huggingface_token,
)

# 모델 로드 함수
def load_pipeline(pipeline_type):
    logging.debug(f'파이프라인 로드 중: {pipeline_type}')
    if pipeline_type == "text2img":
        return StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)
    elif pipeline_type == "img2img":
        return StableDiffusionImg2ImgPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)

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

# 번역 파이프라인 설정
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"

# 디스코드 봇 클래스
class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False
        self.text2img_pipeline = load_pipeline("text2img").to(device)
        self.text2img_pipeline.enable_attention_slicing()  # 메모리 최적화

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        subprocess.Popen(["python", "web.py"])
        logging.info("web.py 서버가 시작되었습니다.")

    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_path = await self.generate_image(prompt_en, negative_prompt)
                user_id = message.author.id  # 사용자의 ID 캡처
                await message.channel.send(
                    f"<@{user_id}> 님이 요청하신 이미지입니다.",
                    file=discord.File(image_path, 'generated_image.png')
                )
            finally:
                self.is_processing = False

    async def generate_image(self, prompt, negative_prompt):
        generator = torch.Generator(device=device).manual_seed(torch.seed())
        images = self.text2img_pipeline(prompt, negative_prompt=negative_prompt, num_inference_steps=50, generator=generator)["images"]
        image_path = f'/tmp/{uuid.uuid4()}.png'
        images[0].save(image_path)
        return image_path

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

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