Update app.py

app.py

@@ -16,24 +16,40 @@ logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %
 hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN"))
 IMAGE_API_URL = "http://211.233.58.201:7896"
 
-def generate_image(prompt: str) -> tuple:
-    """이미지 생성 함수"""
+async def generate_image_prompt(text: str) -> str:
+    """입력 텍스트로부터 이미지 생성용 프롬프트 생성"""
+    try:
+        prompt_messages = [
+            {"role": "system", "content": "You are an expert at creating image generation prompts. Convert the input text into a detailed image generation prompt that describes the visual elements, style, and atmosphere."},
+            {"role": "user", "content": f"Create an image generation prompt for: {text}"}
+        ]
+        
+        response = hf_client.chat_completion(prompt_messages, max_tokens=200)
+        image_prompt = response.choices[0].message.content
+        return f"fantasy style, {image_prompt}"
+    except Exception as e:
+        logging.error(f"Image prompt generation failed: {str(e)}")
+        return f"fantasy style, {text}"
+
+async def generate_image_async(prompt: str) -> tuple:
+    """비동기 이미지 생성 함수"""
     try:
         client = Client(IMAGE_API_URL)
-        # 프롬프트 앞에 "fantasy style," 추가
-        enhanced_prompt = f"fantasy style, {prompt}"
-        result = client.predict(
-            prompt=enhanced_prompt,
-            width=768,
-            height=768,
-            guidance=7.5,
-            inference_steps=30,
-            seed=3,
-            do_img2img=False,
-            init_image=None,
-            image2image_strength=0.8,
-            resize_img=True,
-            api_name="/generate_image"
+        result = await asyncio.get_event_loop().run_in_executor(
+            ThreadPoolExecutor(),
+            lambda: client.predict(
+                prompt=prompt,
+                width=768,
+                height=768,
+                guidance=7.5,
+                inference_steps=30,
+                seed=3,
+                do_img2img=False,
+                init_image=None,
+                image2image_strength=0.8,
+                resize_img=True,
+                api_name="/generate_image"
+            )
         )
         return result[0], result[1]
     except Exception as e:
@@ -49,6 +65,7 @@ def respond(
     top_p=0.9,  
 ):
     system_prefix = """
+    You are 'FantasyAI✨', an advanced AI storyteller specialized in creating immersive fantasy narratives.
     [시스템 프롬프트 내용...]
     """
 
@@ -62,7 +79,6 @@ def respond(
 
     response = ""
     try:
-        # 텍스트 생성
         for msg in hf_client.chat_completion(
             messages,
             max_tokens=max_tokens,
@@ -73,23 +89,27 @@ def respond(
             token = msg.choices[0].delta.content
             if token is not None:
                 response += token.strip("")
-            # 채팅 히스토리에 맞는 형식으로 반환
             history = history + [(message, response)]
-            yield history, None
-
-        # 텍스트 생성이 완료된 후 이미지 생성
-        try:
-            image, seed = generate_image(response[:200])
-            history = history[:-1] + [(message, response)]  # 마지막 메시지 업데이트
-            yield history, image
-        except Exception as e:
-            logging.error(f"Image generation failed: {str(e)}")
-            yield history, None
+            yield history
 
     except Exception as e:
         error_message = f"Error: {str(e)}"
         history = history + [(message, error_message)]
-        yield history, None
+        yield history
+
+async def process_input(message, history, system_message, max_tokens, temperature, top_p):
+    """입력 처리 및 이미지 생성 통합 함수"""
+    # 이미지 프롬프트 생성 및 이미지 생성 시작
+    image_prompt = await generate_image_prompt(message)
+    image_future = asyncio.create_task(generate_image_async(image_prompt))
+    
+    # 텍스트 응답 생성
+    text_response = respond(message, history, system_message, max_tokens, temperature, top_p)
+    
+    # 이미지 생성 완료 대기
+    image, seed = await image_future
+    
+    return text_response, image
 
 # Gradio 인터페이스 설정
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as interface:
@@ -103,11 +123,14 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as interface:
                 label="Chat History",
                 height=500
             )
-            msg = gr.Textbox(
-                label="Enter your message",
-                placeholder="Type your message here...",
-                lines=2
-            )
+            with gr.Row():
+                msg = gr.Textbox(
+                    label="Enter your message",
+                    placeholder="Type your message here...",
+                    lines=2
+                )
+                submit_btn = gr.Button("Submit", variant="primary")
+            
             system_msg = gr.Textbox(
                 label="System Message",
                 value="Write(output) in 한국어.",
@@ -156,10 +179,16 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as interface:
     )
 
     # 이벤트 핸들러
+    submit_btn.click(
+        fn=process_input,
+        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p],
+        outputs=[chatbot, image_output]
+    )
+    
     msg.submit(
-        respond,
-        [msg, chatbot, system_msg, max_tokens, temperature, top_p],
-        [chatbot, image_output]
+        fn=process_input,
+        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p],
+        outputs=[chatbot, image_output]
     )
 
 # 애플리케이션 실행