Update app.py

app.py

@@ -10,6 +10,11 @@ from kolors.models import unet_2d_condition
 from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 import gradio as gr
 import numpy as np
+from huggingface_hub import InferenceClient
+import os
+
+# Cohere 모델 초기화
+client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))
 
 device = "cuda"
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
@@ -53,6 +58,24 @@ pipe_i2i.load_ip_adapter(f'{ckpt_IPA_dir}' , subfolder="", weight_name=["ip_adap
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def call_api(content, system_message, max_tokens=1000, temperature=0.7, top_p=0.95):
+    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
+    response = client.chat_completion(messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
+    return response.choices[0].message['content']
+
+def generate_prompt(korean_prompt):
+    system_message = """
+    Given the following description in Korean, 
+    translate and generate a concise English prompt suitable for a Stable Diffusion model. 
+    The prompt should be focused, descriptive, 
+    and contain specific keywords or phrases that will help guide the image generation process. 
+    Use simple and descriptive language, avoiding unnecessary words. 
+    Ensure the output is in English and follows the format typically used in Stable Diffusion prompts. 
+    The description is: [Insert Korean description here]
+    """
+    optimized_prompt = call_api(korean_prompt, system_message)
+    return optimized_prompt  # 최적화된 프롬프트 반환
+
 @spaces.GPU
 def infer(prompt, 
           ip_adapter_image = None, 
@@ -80,14 +103,15 @@ def infer(prompt,
             height = height,
             generator = generator
         ).images[0] 
-        return image
+        image.save("generated_image.jpg")  # 파일 확장자를 .jpg로 변경
+        return image, "generated_image.jpg"
     else:
         pipe_i2i.to(device)
         image_encoder.to(device)
         pipe_i2i.image_encoder = image_encoder
         pipe_i2i.set_ip_adapter_scale([ip_adapter_scale])
         image = pipe_i2i(
-                prompt=prompt ,
+                prompt=prompt,
                 ip_adapter_image=[ip_adapter_image],
                 negative_prompt=negative_prompt, 
                 height=height,
@@ -97,14 +121,8 @@ def infer(prompt,
                 num_images_per_prompt=1,
                 generator=generator
             ).images[0]
-        return image
-
-examples = [
-    ["一张瓢虫的照片，微距，变焦，高质量，电影，拿着一个牌子，写着“可图”", None, None],
-    ["3D anime style, hyperrealistic oil painting, dolphin leaping out of the water", None, None],
-    ["穿着黑色T恤衫，上面中文绿色大字写着“可图”", "image/test_ip.jpg", 0.5],
-    ["A cute dog is running", "image/test_ip2.png", 0.5]
-]
+        image.save("generated_image.jpg")  # 파일 확장자를 .jpg로 변경
+        return image, "generated_image.jpg"
 
 css="""
 #col-left {
@@ -117,24 +135,33 @@ css="""
 }
 """
 
-def load_description(fp):
-    with open(fp, 'r', encoding='utf-8') as f:
-        content = f.read()
-    return content
-
 with gr.Blocks(css=css) as Kolors:
-    gr.HTML(load_description("assets/title.md"))
     with gr.Row():
         with gr.Column(elem_id="col-left"):
             with gr.Row():
-                prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="Enter your prompt",
+                korean_prompt = gr.Textbox(
+                    label="한국어 프롬프트 입력",
+                    placeholder="한국어로 원하는 프롬프트를 입력하세요",
+                    lines=2
+                )
+            with gr.Row():
+                generate_prompt_button = gr.Button("Generate Prompt")
+            with gr.Row():
+                optimized_prompt = gr.Textbox(
+                    label="최적화된 프롬프트 생성",
+                    placeholder=" ",
+                    lines=2,
+                    interactive=False
+                )
+            with gr.Row():
+                generated_prompt = gr.Textbox(
+                    label="프롬프트 입력",
+                    placeholder="이미지 생성에 사용할 프롬프트를 입력하세요",
                     lines=2
                 )
             with gr.Row():
                 ip_adapter_image = gr.Image(label="Image Prompt (optional)", type="pil")
-            with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row(visible=False):  # Advanced Settings 숨김
                 negative_prompt = gr.Textbox(
                     label="Negative prompt",
                     placeholder="Enter a negative prompt",
@@ -188,23 +215,24 @@ with gr.Blocks(css=css) as Kolors:
                         value=0.5,
                     )
             with gr.Row():
-                run_button = gr.Button("Run")
+                run_button = gr.Button("Generate Image")
             
         with gr.Column(elem_id="col-right"):
             result = gr.Image(label="Result", show_label=False)
-    
-    with gr.Row():
-        gr.Examples(
-                fn = infer,
-                examples = examples,
-                inputs = [prompt, ip_adapter_image, ip_adapter_scale],
-                outputs = [result]
-            )
+            download_button = gr.File(label="Download Image")
 
+    # 최적화된 프롬프트 생성 및 결과 표시
+    generate_prompt_button.click(
+        fn=generate_prompt,
+        inputs=[korean_prompt],
+        outputs=[optimized_prompt]
+    )
+    
+    # 이미지 생성 및 다운로드 파일 경로 설정
     run_button.click(
-        fn = infer,
-        inputs = [prompt, ip_adapter_image, ip_adapter_scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+        fn=infer,
+        inputs=[generated_prompt, ip_adapter_image, ip_adapter_scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, download_button]
     )
 
 Kolors.queue().launch(debug=True)