Update app.py

app.py

@@ -25,8 +25,39 @@ huggingface_token = os.getenv("HF_TOKEN")
 
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
 
+# Hugging Face 클라이언트 초기화
+client = InferenceClient(
+    model="CohereForAI/c4ai-command-r-plus-08-2024",
+    token=huggingface_token
+)
 
+def augment_prompt(prompt):
+    try:
+        if not prompt:
+            return "", gr.Warning("Please enter a prompt first")
+            
+        system_prompt = """You are an expert at writing detailed image generation prompts. 
+        Enhance the given prompt by adding more descriptive details, artistic style, and technical aspects 
+        that would help in generating better images. Keep the core meaning but make it more comprehensive."""
+        
+        full_prompt = f"{system_prompt}\nOriginal prompt: {prompt}\nEnhanced prompt:"
         
+        # 클라이언트 인퍼런스 방식으로 호출
+        response = client.text_generation(
+            full_prompt,
+            max_new_tokens=300,
+            temperature=0.7,
+            top_p=0.95,
+            repetition_penalty=1.1,
+            do_sample=True
+        )
+        
+        enhanced_prompt = response.strip()
+        return enhanced_prompt
+    except Exception as e:
+        print(f"Error in prompt augmentation: {str(e)}")
+        return prompt
+
 #Load prompts for randomization
 df = pd.read_csv('prompts.csv', header=None)
 prompt_values = df.values.flatten()
@@ -826,6 +857,38 @@ input:focus, textarea:focus {
     margin: 0 !important;  /* auto에서 0으로 변경 */
     margin-left: 20px !important;  /* 왼쪽 여백 추가 */
 }
+
+.enhance-button, .generate-button {
+    flex: 1 !important;
+    min-width: 120px !important;
+    height: 40px !important;
+    margin: 0 5px !important;
+}
+
+.enhance-button {
+    background-color: #6b7280 !important;
+    color: white !important;
+}
+
+.enhance-button:hover {
+    background-color: #4b5563 !important;
+}
+.generate-button {
+    background-color: #3b82f6 !important;
+    color: white !important;
+}
+
+.generate-button:hover {
+    background-color: #2563eb !important;
+}
+
+/* 버튼 컨테이너 스타일링 */
+.button-container {
+    display: flex !important;
+    justify-content: space-between !important;
+    gap: 10px !important;
+    width: 100% !important;
+}
 '''
 
 with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
@@ -857,11 +920,30 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
 
     with gr.Tab(label="Generate"):
         # Prompt and Generate Button
+# Prompt and Generate Buttons 영역 수정
         with gr.Row():
             with gr.Column(scale=3):
-                prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
+                prompt = gr.Textbox(
+                    label="Prompt", 
+                    lines=1, 
+                    placeholder="Type a prompt after selecting a LoRA"
+                )
             with gr.Column(scale=1):
-                generate_button = gr.Button("Generate", variant="primary", elem_classes=["button_total"])
+                with gr.Row():
+                    augment_button = gr.Button(
+                        "✨프롬프트 증강", 
+                        variant="secondary", 
+                        size="sm", 
+                        scale=1,
+                        elem_classes=["enhance-button"]
+                    )
+                    generate_button = gr.Button(
+                        "🎨생성", 
+                        variant="primary", 
+                        size="sm", 
+                        scale=1,
+                        elem_classes=["generate-button"]
+                    )
 
         # LoRA Selection Area
         with gr.Row(elem_id="loaded_loras"):
@@ -1009,13 +1091,21 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
                 selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, 
                 lora_image_1, lora_image_2, lora_image_3]
         )
+# 기존 이벤트 핸들러들 위에 추가
+        augment_button.click(
+            fn=augment_prompt,
+            inputs=[prompt],
+            outputs=[prompt],
+            api_name="enhance_prompt"
+        )
 
+# 기존 generate 이벤트 핸들러 유지
         gr.on(
             triggers=[generate_button.click, prompt.submit],
             fn=run_lora,
             inputs=[prompt, input_image, image_strength, cfg_scale, steps, 
-                selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, 
-                randomize_seed, seed, width, height, loras_state],
+                    selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, 
+                    randomize_seed, seed, width, height, loras_state],
             outputs=[result, seed, progress_bar]
         ).then(
             fn=lambda x, history: update_history(x, history) if x is not None else history,
@@ -1023,6 +1113,7 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
             outputs=history_gallery
         )
 
+
 if __name__ == "__main__":
     app.queue(max_size=20)
     app.launch(debug=True)