Update app.py

app.py

@@ -7,14 +7,48 @@ import os
 # Get the Hugging Face token from the environment variable, or a secret if available.
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Check if HF_TOKEN is set; if not, raise a configuration error (handled later)
+# Check if HF_TOKEN is set
 if not HF_TOKEN:
     HF_TOKEN_ERROR = "Hugging Face API token (HF_TOKEN) not found.  Please set it as an environment variable or Gradio secret."
 else:
-    HF_TOKEN_ERROR = None  # No error if the token is found
+    HF_TOKEN_ERROR = None
 
+client = InferenceClient(token=HF_TOKEN)
+PROMPT_IMPROVER_MODEL = "TheBloke/zephyr-7B-beta-AWQ"  # A good general-purpose text model.  AWQ for speed.
+
+def improve_prompt(original_prompt):
+    """Improves the user's prompt using a language model."""
+    if HF_TOKEN_ERROR:
+        raise gr.Error(HF_TOKEN_ERROR)
+
+    try:
+        # Construct a prompt for the language model.
+        system_prompt = "You are a helpful assistant that improves text prompts for image generation models.  Make the prompt more descriptive, detailed, and artistic, while keeping the user's original intent."
+        prompt_for_llm = f"""<|system|>
+{system_prompt}</s>
+<|user|>
+Improve this prompt: {original_prompt}
+</s>
+<|assistant|>
+"""
+        improved_prompt = client.text_generation(
+            prompt=prompt_for_llm,
+            model=PROMPT_IMPROVER_MODEL,
+            max_new_tokens=128,  # Limit the length of the improved prompt
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.2, # Encourage diverse output
+            stop_sequences=["</s>"], # stop at end of sentence
+
+        )
+
+        return improved_prompt.strip()  # Remove leading/trailing whitespace
+
+
+    except Exception as e:
+        print(f"Error improving prompt: {e}")  # Log the error for debugging
+        return original_prompt  # Return the original prompt if there's an error
 
-client = InferenceClient(token=HF_TOKEN)  # Use token instead of provider and api_key
 
 def generate_image(prompt, progress=gr.Progress()):
     """Generates an image using the InferenceClient and provides progress updates."""
@@ -22,30 +56,31 @@ def generate_image(prompt, progress=gr.Progress()):
     if HF_TOKEN_ERROR:
         raise gr.Error(HF_TOKEN_ERROR)
 
-    progress(0, desc="Sending request to Hugging Face...")
+    progress(0, desc="Improving prompt...")
+    improved_prompt = improve_prompt(prompt)
+
+    progress(0.2, desc="Sending request to Hugging Face...") # More granular progress
     try:
-        # Use the client.text_to_image method.  Assume xylaria-iris is valid here.
-        image = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-schnell")
+        image = client.text_to_image(improved_prompt, model="black-forest-labs/FLUX.1-schnell")
 
-        if not isinstance(image, Image.Image): # Basic type checking.
+        if not isinstance(image, Image.Image):
              raise Exception(f"Expected a PIL Image, but got: {type(image)}")
 
         progress(0.8, desc="Processing image...")
-        time.sleep(0.5)  # Simulate some processing
+        time.sleep(0.5)
         progress(1.0, desc="Done!")
-        return image
-    except Exception as e:  # Catch all exceptions from the API call
-        # Check for rate limit errors (different with InferenceClient).  This is a best-effort check.
-        if "rate limit" in str(e).lower(): # Check message, case-insensitively.
+        return image, improved_prompt  # Return both image and improved prompt
+    except Exception as e:
+        if "rate limit" in str(e).lower():
             error_message = f"Rate limit exceeded. Please try again later.  Error: {e}"
         else:
-            error_message = f"An error occurred: {e}"  # Generic error message
+            error_message = f"An error occurred: {e}"
         raise gr.Error(error_message)
 
 
 
-# Gradio Interface (same CSS as before, for consistency)
 css = """
+/* ... (Rest of your CSS, unchanged, from the previous response) ... */
 .container {
     max-width: 800px;
     margin: auto;
@@ -70,14 +105,23 @@ css = """
 .input-section, .output-section {
     margin-bottom: 1.5em;
 }
-
+.output-section img {
+    display: block; /* Ensure image takes full width of container */
+    margin: auto;  /* Center the image horizontally */
+    max-width: 100%; /* Prevent image overflow */
+    height: auto;  /* Maintain aspect ratio */
+    border-radius: 8px;  /* Rounded corners for the image */
+    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);  /* Subtle shadow */
+}
 
 /* Animation for the image appearance - subtle fade-in */
 @keyframes fadeIn {
     from { opacity: 0; transform: translateY(20px); }
     to   { opacity: 1; transform: translateY(0); }
 }
-
+.output-section.animate img {
+    animation: fadeIn 0.8s ease-out;
+}
 
 /* Improve button style */
 .submit-button {
@@ -108,6 +152,16 @@ label{
     display: block;      /* Each label on its own line */
     margin-bottom: 0.5em; /* Space between label and input */
 }
+
+.improved-prompt-display {
+    margin-top: 10px;
+    padding: 8px;
+    border: 1px solid #ccc;
+    border-radius: 4px;
+    background-color: #f9f9f9;
+    font-style: italic;
+    color: #444;
+}
 """
 
 
@@ -115,37 +169,35 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown(
         """
         # Xylaria Iris Image Generator
-        Enter a text prompt and generate an image using the Xylaria Iris model!
+        Enter a text prompt, and we'll enhance it before generating an image!
         """,
         elem_classes="title"
     )
 
-
     with gr.Row():
         with gr.Column():
             with gr.Group(elem_classes="input-section"):
-                prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., A beautiful landscape with a magical tree", lines=3)
+                prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., A cat", lines=3)
                 generate_button = gr.Button("Generate Image", elem_classes="submit-button")
         with gr.Column():
             with gr.Group(elem_classes="output-section") as output_group:
-                image_output = gr.Image(label="Generated Image") # Removed width and height
+                image_output = gr.Image(label="Generated Image", show_download_button=False, interactive=False)  # No SVG, not interactive
+                improved_prompt_output = gr.Textbox(label="Improved Prompt", interactive=False, elem_classes="improved-prompt-display")
 
 
     def on_generate_click(prompt):
         output_group.elem_classes = ["output-section", "animate"]
-        image = generate_image(prompt)
+        image, improved_prompt = generate_image(prompt)
         output_group.elem_classes = ["output-section"]
-        return image
-
+        return image, improved_prompt
 
-    generate_button.click(on_generate_click, inputs=prompt_input, outputs=image_output)
-    prompt_input.submit(on_generate_click, inputs=prompt_input, outputs=image_output)
+    generate_button.click(on_generate_click, inputs=prompt_input, outputs=[image_output, improved_prompt_output])
+    prompt_input.submit(on_generate_click, inputs=prompt_input, outputs=[image_output, improved_prompt_output])
 
     gr.Examples(
-        [["A futuristic cityscape at night"],
-         ["A mystical forest with glowing mushrooms"],
-         ["An astronaut exploring a new planet"],
-         ["A cat wearing a top hat"]],
+        [["A dog"],
+         ["A house on a hill"],
+         ["A spaceship"]],
         inputs=prompt_input
     )