Update app.py

app.py

@@ -18,17 +18,18 @@ from transformers import pipeline
 from gradio_imageslider import ImageSlider
 import numpy as np
 import warnings
-from huggingface_hub import InferenceClient
+import anthropic  # 이 줄 추가
+
+
 
 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
+# Claude 클라이언트 초기화로 대체
+claude_client = anthropic.Anthropic(
+    api_key=os.getenv("CLAUDE_TOKEN")
 )
 
 def augment_prompt(prompt):
@@ -38,21 +39,28 @@ def augment_prompt(prompt):
             
         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."""
+        that would help in generating better images. Keep the core meaning but make it more comprehensive.
+        Focus on visual details and artistic elements that can be represented in an image."""
         
-        full_prompt = f"{system_prompt}\nOriginal prompt: {prompt}\nEnhanced prompt:"
-        
-        # 클라이언트 인퍼런스 방식으로 호출
-        response = client.text_generation(
-            full_prompt,
-            max_new_tokens=300,
+        message = claude_client.messages.create(
+            model="claude-3-sonnet-20240229",
+            max_tokens=300,
             temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.1,
-            do_sample=True
+            messages=[
+                {
+                    "role": "system",
+                    "content": system_prompt
+                },
+                {
+                    "role": "user",
+                    "content": f"Original prompt: {prompt}\nPlease enhance this prompt for image generation."
+                }
+            ]
         )
         
-        enhanced_prompt = response.strip()
+        enhanced_prompt = message.content[0].text.strip()
+        print(f"Original prompt: {prompt}")
+        print(f"Enhanced prompt: {enhanced_prompt}")
         return enhanced_prompt
     except Exception as e:
         print(f"Error in prompt augmentation: {str(e)}")