Update app.py

app.py

@@ -316,25 +316,33 @@ class Demo:
     async def generation_code(self, query: Optional[str], _setting: Dict[str, str]):
         if not query or query.strip() == '':
             query = get_random_placeholder()
-    
+        
+        # 이미지 생성이 필요한지 확인
+        needs_image = '이미지' in query or '그림' in query or 'image' in query.lower()
+        image_prompt = None
+        
+        # 이미지 프롬프트 추출
+        if needs_image:
+            for keyword in ['이미지:', '그림:', 'image:']:
+                if keyword in query.lower():
+                    image_prompt = query.split(keyword)[1].strip()
+                    break
+            if not image_prompt:
+                image_prompt = query  # 명시적 프롬프트가 없으면 전체 쿼리 사용
+        
         messages = [{'role': Role.SYSTEM, 'content': _setting['system']}]
         messages.append({'role': Role.USER, 'content': query})
-    
+        
         system_message = messages[0]['content']
-    
-        claude_messages = [
-            {"role": "user", "content": query}
-        ]
-    
+        claude_messages = [{"role": "user", "content": query}]
         openai_messages = [
             {"role": "system", "content": system_message},
             {"role": "user", "content": query}
         ]
         
         try:
-            # 로딩 시작
             yield [
-                "",  # 빈 문자열로 변경 (코드 출력하지 않음)
+                "",
                 None,
                 gr.update(active_key="loading"),
                 gr.update(open=True)
@@ -345,7 +353,7 @@ class Demo:
             try:
                 async for content in try_claude_api(system_message, claude_messages):
                     yield [
-                        "",  # 빈 문자열로 변경 (코드 출력하지 않음)
+                        "",
                         None,
                         gr.update(active_key="loading"),
                         gr.update(open=True)
@@ -358,7 +366,7 @@ class Demo:
                 
                 async for content in try_openai_api(openai_messages):
                     yield [
-                        "",  # 빈 문자열로 변경 (코드 출력하지 않음)
+                        "",
                         None,
                         gr.update(active_key="loading"),
                         gr.update(open=True)
@@ -367,14 +375,72 @@ class Demo:
                     collected_content = content
 
             if collected_content:
-                # 최종 결과만 표시
+                # 이미지 생성이 필요한 경우
+                if needs_image and image_prompt:
+                    try:
+                        print(f"Generating image for prompt: {image_prompt}")
+                        # FLUX 모델을 사용하여 이미지 생성
+                        if pipe is not None:
+                            image = generate_image(
+                                prompt=image_prompt,
+                                height=512,
+                                width=512,
+                                steps=8,
+                                scales=3.5,
+                                seed=random.randint(1, 10000)
+                            )
+                            
+                            # 이미지를 Base64로 인코딩
+                            buffered = BytesIO()
+                            image.save(buffered, format="PNG")
+                            img_str = base64.b64encode(buffered.getvalue()).decode()
+                            
+                            # HTML에 이미지 추가
+                            image_html = f'''
+                            <div class="generated-image" style="margin: 20px 0; text-align: center;">
+                                <h3 style="color: #333; margin-bottom: 10px;">Generated Image:</h3>
+                                <img src="data:image/png;base64,{img_str}" 
+                                     style="max-width: 100%; 
+                                            border-radius: 10px; 
+                                            box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
+                                <p style="color: #666; margin-top: 10px; font-style: italic;">
+                                    Prompt: {html.escape(image_prompt)}
+                                </p>
+                            </div>
+                            '''
+                            
+                            # HTML 응답에 이미지 삽입
+                            if '```html' in collected_content:
+                                # HTML 코드 블록 내부에 이미지 추가
+                                collected_content = collected_content.replace('```html\n', f'```html\n{image_html}')
+                            else:
+                                # HTML 코드 블록으로 감싸서 이미지 추가
+                                collected_content = f'```html\n{image_html}\n```\n{collected_content}'
+                            
+                            print("Image generation successful")
+                        else:
+                            raise Exception("FLUX model not initialized")
+                            
+                    except Exception as e:
+                        print(f"Image generation error: {str(e)}")
+                        error_message = f'''
+                        <div style="color: #ff4d4f; padding: 10px; margin: 10px 0; 
+                                    border-left: 4px solid #ff4d4f; background: #fff2f0;">
+                            <p>Failed to generate image: {str(e)}</p>
+                        </div>
+                        '''
+                        if '```html' in collected_content:
+                            collected_content = collected_content.replace('```html\n', f'```html\n{error_message}')
+                        else:
+                            collected_content = f'```html\n{error_message}\n```\n{collected_content}'
+
+                # 최종 결과 표시
                 yield [
                     collected_content,
                     send_to_sandbox(remove_code_block(collected_content)),
                     gr.update(active_key="render"),
                     gr.update(open=False)
                 ]
-
             else:
                 raise ValueError("No content was generated from either API")
                 
@@ -382,6 +448,9 @@ class Demo:
             print(f"Error details: {str(e)}")
             raise ValueError(f'Error calling APIs: {str(e)}')
 
+    def clear_history(self):
+        return []
+
 def remove_code_block(text):
     pattern = r'```html\n(.+?)\n```'
     match = re.search(pattern, text, re.DOTALL)