Create app.py

app.py

@@ -0,0 +1,244 @@
+import gradio as gr
+import replicate
+import requests
+import os
+import json
+from io import BytesIO
+from PIL import Image
+
+# 환경 변수에서 토큰 가져오기
+REPLICATE_API_TOKEN = os.getenv("RAPI_TOKEN")
+FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN")
+
+def generate_prompt_with_llm(topic):
+    """주제를 받아서 LLM을 사용해 이미지 프롬프트를 생성"""
+    url = "https://api.friendli.ai/dedicated/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {FRIENDLI_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    
+    # 프롬프트 생성을 위한 시스템 메시지
+    system_prompt = """당신은 이미지 생성을 위한 프롬프트를 만드는 전문가입니다. 
+    사용자가 주제를 제공하면, 그 주제에 맞는 상세하고 창의적인 이미지 프롬프트를 생성해주세요.
+    프롬프트는 구체적인 시각적 요소, 스타일, 분위기, 색상 등을 포함해야 합니다.
+    응답은 프롬프트만 반환하고 다른 설명은 포함하지 마세요."""
+    
+    payload = {
+        "model": "dep89a2fld32mcm",
+        "messages": [
+            {
+                "role": "system",
+                "content": system_prompt
+            },
+            {
+                "role": "user",
+                "content": f"다음 주제에 대한 이미지 프롬프트를 생성해주세요: {topic}"
+            }
+        ],
+        "max_tokens": 200,
+        "top_p": 0.8,
+        "stream": False
+    }
+    
+    try:
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code == 200:
+            result = response.json()
+            return result['choices'][0]['message']['content'].strip()
+        else:
+            return f"프롬프트 생성 실패: {response.status_code}"
+    except Exception as e:
+        return f"프롬프트 생성 중 오류 발생: {str(e)}"
+
+def translate_to_english(text):
+    """한글 텍스트를 영어로 번역 (LLM 사용)"""
+    # 간단한 한글 체크
+    if not any(ord('가') <= ord(char) <= ord('힣') for char in text):
+        return text  # 한글이 없으면 그대로 반환
+    
+    url = "https://api.friendli.ai/dedicated/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {FRIENDLI_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    
+    payload = {
+        "model": "dep89a2fld32mcm",
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are a translator. Translate the given Korean text to English. Only return the translation without any explanation."
+            },
+            {
+                "role": "user",
+                "content": text
+            }
+        ],
+        "max_tokens": 500,
+        "top_p": 0.8,
+        "stream": False
+    }
+    
+    try:
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code == 200:
+            result = response.json()
+            return result['choices'][0]['message']['content'].strip()
+        else:
+            return text  # 번역 실패시 원본 반환
+    except Exception as e:
+        return text  # 오류 발생시 원본 반환
+
+def generate_image(prompt, seed=10):
+    """Replicate API를 사용해 이미지 생성"""
+    try:
+        # 한글이 포함된 경우 영어로 번역
+        english_prompt = translate_to_english(prompt)
+        
+        # Replicate 클라이언트 초기화
+        if not REPLICATE_API_TOKEN:
+            return None, "RAPI_TOKEN 환경변수가 설정되지 않았습니다."
+        
+        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
+        
+        # 이미지 생성
+        input_params = {
+            "seed": seed,
+            "prompt": english_prompt,
+            "speed_mode": "Extra Juiced 🚀 (even more speed)",
+            "output_quality": 80
+        }
+        
+        output = client.run(
+            "prunaai/hidream-l1-fast:17c237d753218fed0ed477cb553902b6b75735f48c128537ab829096ef3d3645",
+            input=input_params
+        )
+        
+        # 이미지 데이터 읽기
+        if output:
+            # output이 URL인 경우
+            if isinstance(output, str) and output.startswith('http'):
+                response = requests.get(output)
+                img = Image.open(BytesIO(response.content))
+                return img, f"사용된 영어 프롬프트: {english_prompt}"
+            # output이 바이너리 데이터인 경우
+            else:
+                img = Image.open(BytesIO(output.read()))
+                return img, f"사용된 영어 프롬프트: {english_prompt}"
+        else:
+            return None, "이미지 생성 실패"
+            
+    except Exception as e:
+        return None, f"이미지 생성 중 오류 발생: {str(e)}"
+
+def process_topic_to_image(topic, manual_prompt, use_manual, seed):
+    """주제를 받아서 프롬프트 생성 후 이미지 생성"""
+    if use_manual and manual_prompt.strip():
+        # 수동 프롬프트 사용
+        prompt = manual_prompt.strip()
+    else:
+        # LLM으로 프롬프트 생성
+        prompt = generate_prompt_with_llm(topic)
+        if "오류" in prompt or "실패" in prompt:
+            return None, prompt, prompt
+    
+    # 이미지 생성
+    image, status = generate_image(prompt, seed)
+    
+    return image, prompt, status
+
+# Gradio 인터페이스 생성
+with gr.Blocks(title="AI 이미지 생성기") as demo:
+    gr.Markdown("""
+    # 🎨 AI 이미지 생성기
+    
+    주제를 입력하면 AI가 프롬프트를 생성하고 이미지를 만들어줍니다.
+    프롬프트는 자동 생성되거나 직접 편집할 수 있으며, 한글 프롬프트는 자동으로 영어로 번역됩니다.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            topic_input = gr.Textbox(
+                label="주제 입력",
+                placeholder="예: 우주를 여행하는 고양이",
+                lines=2
+            )
+            
+            generated_prompt = gr.Textbox(
+                label="생성된 프롬프트 (편집 가능)",
+                placeholder="자동으로 생성된 프롬프트가 여기에 표시됩니다.",
+                lines=3
+            )
+            
+            use_manual_checkbox = gr.Checkbox(
+                label="수동 프롬프트 사용",
+                value=False
+            )
+            
+            seed_input = gr.Slider(
+                minimum=1,
+                maximum=100,
+                value=10,
+                step=1,
+                label="시드 값 (동일한 시드는 동일한 이미지 생성)"
+            )
+            
+            generate_prompt_btn = gr.Button("프롬프트 생성", variant="secondary")
+            generate_image_btn = gr.Button("이미지 생성", variant="primary")
+            
+        with gr.Column(scale=1):
+            output_image = gr.Image(
+                label="생성된 이미지",
+                type="pil"
+            )
+            
+            status_text = gr.Textbox(
+                label="상태 메시지",
+                lines=2,
+                interactive=False
+            )
+    
+    # 프롬프트 생성 버튼 클릭 이벤트
+    def generate_prompt_only(topic):
+        if not topic.strip():
+            return "주제를 입력해주세요."
+        prompt = generate_prompt_with_llm(topic)
+        return prompt
+    
+    generate_prompt_btn.click(
+        fn=generate_prompt_only,
+        inputs=[topic_input],
+        outputs=[generated_prompt]
+    )
+    
+    # 이미지 생성 버튼 클릭 이벤트
+    generate_image_btn.click(
+        fn=process_topic_to_image,
+        inputs=[topic_input, generated_prompt, use_manual_checkbox, seed_input],
+        outputs=[output_image, generated_prompt, status_text]
+    )
+    
+    gr.Markdown("""
+    ---
+    ### 사용 방법:
+    1. **주제 입력**: 원하는 이미지의 주제를 입력하세요.
+    2. **프롬프트 생성**: "프롬프트 생성" 버튼을 클릭하면 AI가 자동으로 프롬프트를 생성합니다.
+    3. **프롬프트 편집**: 생성된 프롬프트를 원하는 대로 수정할 수 있습니다.
+    4. **이미지 생성**: "이미지 생성" 버튼을 클릭하면 이미지가 생성됩니다.
+    
+    💡 **팁**: 
+    - 한글 프롬프트는 자동으로 영어로 번역되어 전송됩니다.
+    - 시드 값을 고정하면 동일한 프롬프트로 동일한 이미지를 재생성할 수 있습니다.
+    - "수동 프롬프트 사용"을 체크하면 주제 기반 자동 생성을 건너뛰고 직접 입력한 프롬프트를 사용합니다.
+    """)
+
+# 앱 실행
+if __name__ == "__main__":
+    # 환경 변수 확인
+    if not REPLICATE_API_TOKEN:
+        print("경고: RAPI_TOKEN 환경 변수가 설정되지 않았습니다.")
+    if not FRIENDLI_TOKEN:
+        print("경고: FRIENDLI_TOKEN 환경 변수가 설정되지 않았습니다.")
+    
+    demo.launch()