Update app.py

app.py

@@ -7,21 +7,28 @@ from huggingface_hub import InferenceClient
 API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment variable
 API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"
 
-enhancer = InferenceClient(api_key =API_TOKEN)
 
-for message in enhancer.chat_completion(
-	model="meta-llama/Llama-3.2-1B-Instruct",
-	messages=[{"role": "user", "content": "What is the capital of France?"}],
-	max_tokens=500,
-	stream=True,
-):
-    print(message.choices[0].delta.content, end="")
-# Load API Token from environment variable
-
-# Hugging Face Inference API URL
+def enhanche_prompt(prompt, system_prompt='You are a prompt enhancer',  model="meta-llama/Llama-3.2-1B-Instruct", max_tokens=512, stream=False ):
+    enhancer = InferenceClient(api_key = API_TOKEN)
+    for message in enhancer.chat_completion(
+    	model=model,
+    	messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt}
+        ],
+    	max_tokens=max_tokens,
+    	stream=stream,
+    ):
+        result = message.choices[0].delta.content
+        print(result, end="")
+        return result
+        
 
 # Function to call Hugging Face API and get the generated image
-def generate_image(prompt):
+def generate_image(prompt, enhance=False]):
+    if enhance:
+        prompt = enhanche_prompt(prompt)
+        
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
     data = {"inputs": prompt}
     
@@ -72,7 +79,7 @@ css = """
 }
 #app-container3 {
     background-color: rgba(255, 255, 255, 0.001);  /* Corrected to make semi-transparent */
-    max-width: 600px;
+    max-width: 300px;
     margin-left: auto;
     margin-right: auto;
     margin-bottom: 10px;
@@ -81,7 +88,7 @@ css = """
 }
 #app-container {
     background-color: rgba(255, 255, 255, 0.001);  /* Semi-transparent background */
-    max-width: 600px;
+    max-width: 300px;
     margin: 0 auto;  /* Center horizontally */
     padding-bottom: 10px;
     border-radius: 25px;
@@ -106,9 +113,6 @@ css = """
 """
 
 
-
-
-
 # Create Gradio interface
 def create_ui():
     with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as ui: