Update app.py

app.py

@@ -8,7 +8,6 @@ import io
 from PIL import Image
 import time
 
-
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
 
@@ -19,6 +18,9 @@ intents.message_content = True
 # 번역 파이프라인 설정
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
+# 환경 변수에서 지정된 채널 ID 가져오기
+TARGET_CHANNEL_ID = os.getenv('TARGET_CHANNEL_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"
 
@@ -31,24 +33,17 @@ def translate_prompt(prompt):
     return translated_text
 
 def generate_image(prompt, negative_prompt):
-    headers = {
-        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
-    }
-    # 프롬프트와 네거티브 프롬프트를 결합하여 하나의 문자열로 처리
+    headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
     combined_prompt = f"{prompt}. {negative_prompt}"
-    data = {
-        "inputs": combined_prompt
-    }
+    data = {"inputs": combined_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
+        return response.json()[0]['url']
     else:
         logging.error(f"API 요청에 실패했습니다: {response.text}")
         return None
 
-
 class MyClient(discord.Client):
     async def on_ready(self):
         logging.info(f'{self.user}로 로그인되었습니다!')
@@ -56,52 +51,26 @@ class MyClient(discord.Client):
         logging.info("web.py 서버가 시작되었습니다.")
 
     async def on_message(self, message):
-        if message.author == self.user:
+        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)
-                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}"
-                    )
+                    await message.channel.send(f"<@{message.author.id}> 님이 요청하신 이미지입니다: {image_url}")
                 else:
-                    await message.channel.send(f"<@{user_id}> 이미지 생성에 실패하였습니다.")
+                    await message.channel.send(f"<@{message.author.id}> 이미지 생성에 실패하였습니다.")
             finally:
                 self.is_processing = False
         else:
-            # "!image" 명령어가 아닐 경우 안내 메시지 전송
             await message.channel.send('올바른 명령어를 입력해 주세요. 예) "!image 귀여운 고양이가 잠을 자고있다." 등으로 입력하시면 이미지가 생성됩니다.')
 
-
-def generate_image_with_retry(prompt, negative_prompt, max_retries=3):
-    headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
-    combined_prompt = f"{prompt}. {negative_prompt}"
-    data = {"inputs": combined_prompt}
-    api_url = "https://api-inference.huggingface.co/models/fluently/Fluently-XL-Final"
-
-    for attempt in range(max_retries):
-        response = requests.post(api_url, headers=headers, json=data)
-        if response.status_code == 200:
-            return response.json()[0]['url']
-        elif response.status_code == 503 and 'estimated_time' in response.json():
-            wait_time = response.json()['estimated_time']
-            logging.info(f"모델 로딩 중, {wait_time}초 후 재시도합니다.")
-            time.sleep(wait_time)  # 대기 후 재시도
-        else:
-            logging.error(f"API 요청에 실패했습니다: {response.text}")
-            return None
-    logging.error("최대 재시도 횟수 초과")
-    return None
 # 봇 실행
 if __name__ == "__main__":
     discord_token = os.getenv('DISCORD_TOKEN')
     discord_client = MyClient(intents=intents)
     discord_client.run(discord_token)
+