Update app.py

app.py

@@ -433,197 +433,201 @@ async def try_openai_api(openai_messages):
 
 class Demo:
     def __init__(self):
-        # file_state 대신 직접 속성으로 정의
-        self.last_analysis = {}  # 파일 분석 결과 저장
-        self.uploaded_files = {}  # 업로드된 파일 객체 저장
+        self.last_analysis = {}
+        self.uploaded_files = {}
 
     async def generation_code(self, query: Optional[str], _setting: Dict[str, str]):
-        if not query or query.strip() == '':
-            query = get_random_placeholder()
-        
-        # 파일 분석 결과가 있는 경우 컨텍스트로 추가
-        context = ""
-        if self.last_analysis:  # file_state.last_analysis 대신 직접 참조
-            context = "Based on the uploaded files:\n"
-            for filename, analysis in self.last_analysis.items():
-                context += f"\nFile '{filename}':\n{analysis}\n"
-            query = f"{context}\n\nUser Query: {query}"
-
-        # 이미지 생성이 필요한지 확인
-        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}]
+        try:
+            if not query or query.strip() == '':
+                query = get_random_placeholder()
+            
+            # 파일 분석 결과가 있는 경우 컨텍스트로 추가
+            context = ""
+            if self.last_analysis:  # self.file_state 대신 직접 self.last_analysis 사용
+                context = "Based on the uploaded files:\n"
+                for filename, analysis in self.last_analysis.items():
+                    context += f"\nFile '{filename}':\n{analysis}\n"
+                query = f"{context}\n\nUser Query: {query}"
+
+            # 이미지 생성이 필요한지 확인
+            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 API 메시지에 추가
-        for filename, file in self.uploaded_files.items():  # file_state.uploaded_files 대신 직접 참조
-            if any(filename.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']):
-                try:
-                    image = Image.open(file)
-                    image_bytes = io.BytesIO()
-                    image.save(image_bytes, format=image.format)
-                    image_base64 = base64.b64encode(image_bytes.getvalue()).decode('utf-8')
-                    
-                    claude_messages[0]["content"] = [
-                        {
-                            "type": "image",
-                            "source": {
-                                "type": "base64",
-                                "media_type": f"image/{image.format.lower()}",
-                                "data": image_base64
+            # 업로드된 이미지가 있는 경우 Claude API 메시지에 추가
+            for filename, file in self.uploaded_files.items():  # self.file_state 대신 직접 self.uploaded_files 사용
+                if any(filename.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']):
+                    try:
+                        image = Image.open(file)
+                        image_bytes = io.BytesIO()
+                        image.save(image_bytes, format=image.format)
+                        image_base64 = base64.b64encode(image_bytes.getvalue()).decode('utf-8')
+                        
+                        claude_messages[0]["content"] = [
+                            {
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": f"image/{image.format.lower()}",
+                                    "data": image_base64
+                                }
+                            },
+                            {
+                                "type": "text",
+                                "text": query
                             }
-                        },
-                        {
-                            "type": "text",
-                            "text": query
-                        }
-                    ]
-                except Exception as e:
-                    print(f"Error processing uploaded image: {str(e)}")
+                        ]
+                    except Exception as e:
+                        print(f"Error processing uploaded image: {str(e)}")
 
-        openai_messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": query}
-        ]
-        
-        try:
-            yield [
-                "",
-                None,
-                gr.update(active_key="loading"),
-                gr.update(open=True)
+            openai_messages = [
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": query}
             ]
-            await asyncio.sleep(0)
-
-            collected_content = None
+            
             try:
-                async for content in try_claude_api(system_message, claude_messages):
-                    yield [
-                        "",
-                        None,
-                        gr.update(active_key="loading"),
-                        gr.update(open=True)
-                    ]
-                    await asyncio.sleep(0)
-                    collected_content = content
-                    
-            except Exception as claude_error:
-                print(f"Falling back to OpenAI API due to Claude error: {str(claude_error)}")
-                
-                async for content in try_openai_api(openai_messages):
-                    yield [
-                        "",
-                        None,
-                        gr.update(active_key="loading"),
-                        gr.update(open=True)
-                    ]
-                    await asyncio.sleep(0)
-                    collected_content = content
+                yield [
+                    "",
+                    None,
+                    gr.update(active_key="loading"),
+                    gr.update(open=True)
+                ]
+                await asyncio.sleep(0)
 
-            if collected_content:
-                # FLUX를 사용한 이미지 생성이 필요한 경우
-                if needs_image and image_prompt:
-                    try:
-                        print(f"Generating image for prompt: {image_prompt}")
-                        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)
-                            )
-                            
-                            buffered = BytesIO()
-                            image.save(buffered, format="PNG")
-                            img_str = base64.b64encode(buffered.getvalue()).decode()
-                            
-                            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>
-                            '''
-                            
-                            if '```html' in collected_content:
-                                collected_content = collected_content.replace('```html\n', f'```html\n{image_html}')
-                            else:
-                                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}'
-
-                # 업로드된 이미지 표시
-                for filename, file in self.file_state.uploaded_files.items():
-                    if any(filename.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']):
+                collected_content = None
+                try:
+                    async for content in try_claude_api(system_message, claude_messages):
+                        yield [
+                            "",
+                            None,
+                            gr.update(active_key="loading"),
+                            gr.update(open=True)
+                        ]
+                        await asyncio.sleep(0)
+                        collected_content = content
+                        
+                except Exception as claude_error:
+                    print(f"Falling back to OpenAI API due to Claude error: {str(claude_error)}")
+                    
+                    async for content in try_openai_api(openai_messages):
+                        yield [
+                            "",
+                            None,
+                            gr.update(active_key="loading"),
+                            gr.update(open=True)
+                        ]
+                        await asyncio.sleep(0)
+                        collected_content = content
+
+                if collected_content:
+                    # 이미지 생성이 필요한 경우
+                    if needs_image and image_prompt:
                         try:
-                            image = Image.open(file)
-                            buffered = BytesIO()
-                            image.save(buffered, format=image.format)
-                            img_str = base64.b64encode(buffered.getvalue()).decode()
-                            
-                            uploaded_image_html = f'''
-                            <div class="uploaded-image" style="margin: 20px 0; text-align: center;">
-                                <h3 style="color: #333; margin-bottom: 10px;">Uploaded Image: {html.escape(filename)}</h3>
-                                <img src="data:image/{image.format.lower()};base64,{img_str}" 
-                                     style="max-width: 100%; 
-                                            border-radius: 10px; 
-                                            box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
+                            print(f"Generating image for prompt: {image_prompt}")
+                            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)
+                                )
+                                
+                                buffered = BytesIO()
+                                image.save(buffered, format="PNG")
+                                img_str = base64.b64encode(buffered.getvalue()).decode()
+                                
+                                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>
+                                '''
+                                
+                                if '```html' in collected_content:
+                                    collected_content = collected_content.replace('```html\n', f'```html\n{image_html}')
+                                else:
+                                    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{uploaded_image_html}')
+                                collected_content = collected_content.replace('```html\n', f'```html\n{error_message}')
                             else:
-                                collected_content = f'```html\n{uploaded_image_html}\n```\n{collected_content}'
-                        except Exception as e:
-                            print(f"Error displaying uploaded image: {str(e)}")
+                                collected_content = f'```html\n{error_message}\n```\n{collected_content}'
+
+                    # 업로드된 이미지 표시
+                    for filename, file in self.uploaded_files.items():
+                        if any(filename.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']):
+                            try:
+                                image = Image.open(file)
+                                buffered = BytesIO()
+                                image.save(buffered, format=image.format)
+                                img_str = base64.b64encode(buffered.getvalue()).decode()
+                                
+                                uploaded_image_html = f'''
+                                <div class="uploaded-image" style="margin: 20px 0; text-align: center;">
+                                    <h3 style="color: #333; margin-bottom: 10px;">Uploaded Image: {html.escape(filename)}</h3>
+                                    <img src="data:image/{image.format.lower()};base64,{img_str}" 
+                                         style="max-width: 100%; 
+                                                border-radius: 10px; 
+                                                box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
+                                </div>
+                                '''
+                                
+                                if '```html' in collected_content:
+                                    collected_content = collected_content.replace('```html\n', f'```html\n{uploaded_image_html}')
+                                else:
+                                    collected_content = f'```html\n{uploaded_image_html}\n```\n{collected_content}'
+                            except Exception as e:
+                                print(f"Error displaying uploaded image: {str(e)}")
+
+                    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")
+                    
+            except Exception as e:
+                print(f"Error details: {str(e)}")
+                raise ValueError(f'Error calling APIs: {str(e)}')
 
-                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")
-                
         except Exception as e:
-            print(f"Error details: {str(e)}")
-            raise ValueError(f'Error calling APIs: {str(e)}')
+            print(f"Error in generation_code: {str(e)}")
+            raise e
 
     async def handle_file_upload(self, files):
         """파일 업로드 처리 함수"""