Update app.py

app.py

@@ -1,82 +1,32 @@
-import gradio as gr
-from gradio_client import Client
+import random
 import os
-from dotenv import load_dotenv
+import uuid
+from datetime import datetime
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
 import warnings
 
 # 경고 메시지 숨기기
 warnings.filterwarnings('ignore', category=UserWarning)
 
-# 환경 변수 로드
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
+# 저장 디렉토리 생성
+SAVE_DIR = "saved_images"
+if not os.path.exists(SAVE_DIR):
+    os.makedirs(SAVE_DIR, exist_ok=True)
 
-def generate_diagram(prompt, width=1024, height=1024):
-    """FLUX AI를 사용하여 다이어그램 생성"""
-    try:
-        # API 호출을 통해 이미지 생성
-        result = gr.Image.update(value=None)  # 초기 상태는 빈 이미지
-        
-        # Hugging Face API 호출
-        client = Client(
-            "https://black-forest-labs-flux-1-schnell.hf.space",  # API 엔드포인트 직접 지정
-            hf_token=HF_TOKEN,
-        )
-        
-        # 이미지 생성
-        result = client.predict(
-            prompt,
-            1872187377,  # seed
-            False,       # randomize_seed
-            width,
-            height,
-            4,          # num_inference_steps
-            api_name="/infer"
-        )
-        return result
-    except Exception as e:
-        raise gr.Error(f"다이어그램 생성 중 오류 발생: {str(e)}")
+# 장치 설정
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Convert example format for Gradio
-GRADIO_EXAMPLES = [
-    [
-        """A handrawn colorful mind map diagram, educational style, vibrant colors, clear hierarchy.
-        KNOWLEDGE
-        ├── ACQUISITION [Brain with Lightning ~60px]
-        │   ├── READING [Open Book with Glow]
-        │   └── PRACTICE [Hands-on Tools]
-        └── APPLICATION
-            ├── CREATION [Artist Palette]
-            └── INNOVATION [Lightbulb]""",
-        1024,
-        1024
-    ],
-    [
-        """A handrawn colorful mind map diagram, tech-focused style, neon accents.
-        DIGITAL TRANSFORM
-        ├── CLOUD [Cloud with Data ~55px]
-        │   ├── STORAGE [Database]
-        │   └── COMPUTING [Server]
-        └── SECURITY
-            ├── ENCRYPTION [Lock]
-            └── MONITORING [Radar]""",
-        1024,
-        1024
-    ],
-    [
-        """A handrawn colorful mind map diagram, creative style, flowing design.
-        INNOVATION
-        ├── IDEATION [Lightbulb ~60px]
-        │   ├── RESEARCH [Magnifier]
-        │   └── BRAINSTORM [Brain]
-        └── EXECUTION
-            ├── PROTOTYPE [Tools]
-            └── TEST [Checklist]""",
-        1024,
-        1024
-    ]
-]
+# 모델 로드
+repo_id = "black-forest-labs/FLUX.1-schnell"
+pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float32)
+pipeline = pipeline.to(device)
 
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
 
 # Enhanced examples with more detailed prompts and specific styling
 EXAMPLES = [
@@ -273,15 +223,41 @@ EXAMPLES = [
     }
 ]
 
-
-
 # Convert examples to Gradio format
 GRADIO_EXAMPLES = [
     [example["prompt"], example["width"], example["height"]]
     for example in EXAMPLES
 ]
 
+def generate_diagram(prompt, width=1024, height=1024):
+    """FLUX AI를 사용하여 다이어그램 생성"""
+    try:
+        # 시드 설정
+        seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device=device).manual_seed(seed)
+        
+        # 이미지 생성
+        image = pipeline(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=4,
+            generator=generator,
+        ).images[0]
+        
+        # 이미지 저장
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        unique_id = str(uuid.uuid4())[:8]
+        filename = f"diagram_{timestamp}_{unique_id}.png"
+        save_path = os.path.join(SAVE_DIR, filename)
+        image.save(save_path)
+        
+        return image
+        
+    except Exception as e:
+        raise gr.Error(f"다이어그램 생성 중 오류 발생: {str(e)}")
 
+# Gradio 인터페이스 생성
 demo = gr.Interface(
     fn=generate_diagram,
     inputs=[
@@ -293,14 +269,14 @@ demo = gr.Interface(
         gr.Slider(
             label="너비",
             minimum=512,
-            maximum=2048,
+            maximum=MAX_IMAGE_SIZE,
             step=128,
             value=1024
         ),
         gr.Slider(
             label="높이",
             minimum=512,
-            maximum=2048,
+            maximum=MAX_IMAGE_SIZE,
             step=128,
             value=1024
         )
@@ -319,9 +295,9 @@ demo = gr.Interface(
     cache_examples=True
 )
 
-# 앱 실행 부분만 수정
+# 앱 실행
 if __name__ == "__main__":
-    demo.queue()  # 큐 기능을 별도로 활성화
+    demo.queue()
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,