Update app.py

app.py

@@ -53,13 +53,16 @@ class MyClient(discord.Client):
 async def generate_image(prompt):
     try:
         # 모델 호출
-        response = hf_client(inputs=prompt)
-        image_data = response['images'][0]  # 응답에서 이미지 데이터 추출
-        image_bytes = io.BytesIO(base64.b64decode(image_data))
-        image = Image.open(image_bytes)
-        image_path = "output.png"
-        image.save(image_path)
-        return image_path
+        response = hf_client(inputs={"inputs": prompt}, options={"wait_for_model": True})  # options는 필요에 따라 조정 가능
+        image_data = response.get('images', [None])[0]  # 응답에서 이미지 데이터 추출
+        if image_data:
+            image_bytes = io.BytesIO(base64.b64decode(image_data))
+            image = Image.open(image_bytes)
+            image_path = "output.png"
+            image.save(image_path)
+            return image_path
+        else:
+            raise ValueError("No image data received.")
     except Exception as e:
         logging.error(f"Failed to generate image: {str(e)}")
         return None