Update app.py

app.py

@@ -433,12 +433,20 @@ async def try_openai_api(openai_messages):
 
 class Demo:
     def __init__(self):
-        pass
+        self.file_state = FileAnalysisState()  # 파일 상태 관리 추가
 
     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.file_state.last_analysis:
+            context = "Based on the uploaded files:\n"
+            for filename, analysis in self.file_state.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
@@ -450,13 +458,40 @@ class Demo:
                     image_prompt = query.split(keyword)[1].strip()
                     break
             if not image_prompt:
-                image_prompt = query  # 명시적 프롬프트가 없으면 전체 쿼리 사용
-        
+                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.file_state.uploaded_files.items():
+            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
+                        }
+                    ]
+                except Exception as e:
+                    print(f"Error processing uploaded image: {str(e)}")
+
         openai_messages = [
             {"role": "system", "content": system_message},
             {"role": "user", "content": query}
@@ -497,11 +532,10 @@ class Demo:
                     collected_content = content
 
             if collected_content:
-                # 이미지 생성이 필요한 경우
+                # FLUX를 사용한 이미지 생성이 필요한 경우
                 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,
@@ -512,12 +546,10 @@ class Demo:
                                 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>
@@ -531,12 +563,9 @@ class Demo:
                             </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")
@@ -556,7 +585,32 @@ class Demo:
                         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']):
+                        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)),
@@ -571,8 +625,15 @@ class Demo:
             raise ValueError(f'Error calling APIs: {str(e)}')
 
     def clear_history(self):
+        self.file_state.last_analysis.clear()
+        self.file_state.uploaded_files.clear()
         return []
 
+    def update_file_state(self, analysis_result, files):
+        """파일 분석 결과와 파일 객체 업데이트"""
+        self.file_state.last_analysis.update(analysis_result)
+        self.file_state.uploaded_files.update(files)
+
 def remove_code_block(text):
     pattern = r'```html\n(.+?)\n```'
     match = re.search(pattern, text, re.DOTALL)
@@ -686,7 +747,13 @@ def get_random_placeholder():
 def update_placeholder():
     return gr.update(placeholder=get_random_placeholder())
 
+# 파일 분석 결과를 저장할 전역 변수 추가
+class FileAnalysisState:
+    def __init__(self):
+        self.last_analysis = {}  # {파일명: 분석결과} 형태로 저장
+        self.uploaded_files = {}  # {파일명: 파일객체} 형태로 저장
 
+file_state = FileAnalysisState()
 
 def create_main_interface():
     """메인 인터페이스 생성 함수"""
@@ -1140,29 +1207,34 @@ Use the "Generate" button for basic creation, "Enhance" button for prompt improv
                                         with antd.Tabs.Item(key="render"):
                                             sandbox = gr.HTML(elem_classes="html_content")
 
-                        # 이벤트 핸들러 정의
                         async def handle_file_upload(files):
                             if not files:
                                 return "No files uploaded"
-                            
+    
                             if len(files) > FileProcessor.MAX_FILES:
                                 return f"Maximum {FileProcessor.MAX_FILES} files allowed"
-                            
+    
                             results = []
+                            file_state.last_analysis.clear()
+                            file_state.uploaded_files.clear()
+    
                             for file in files:
                                 if not FileProcessor.is_allowed_file(file.name):
                                     results.append(f"Unsupported file: {file.name}")
                                     continue
-                                
+        
                                 if os.path.getsize(file.name) > FileProcessor.MAX_FILE_SIZE:
                                     results.append(f"File too large: {file.name}")
                                     continue
-                                
+        
                                 result = await FileProcessor.process_file(file)
+                                file_state.last_analysis[file.name] = result
+                                file_state.uploaded_files[file.name] = file
                                 results.append(f"Analysis for {file.name}:\n{result}\n")
-                            
+    
                             return "\n".join(results)
 
+
                         # 이벤트 핸들러 연결
                         analyze_btn.click(
                             fn=handle_file_upload,