Update app.py

app.py

@@ -3,26 +3,24 @@ from huggingface_hub import InferenceClient
 from PIL import Image
 import time
 import os
+import base64
+from io import BytesIO
 
-# 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 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
 
 client = InferenceClient(token=HF_TOKEN)
-PROMPT_IMPROVER_MODEL = "TheBloke/zephyr-7B-beta-AWQ"  # A good general-purpose text model.  AWQ for speed.
+PROMPT_IMPROVER_MODEL = "HuggingFaceH4/zephyr-7b-beta"
 
 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>
@@ -34,32 +32,29 @@ Improve this prompt: {original_prompt}
         improved_prompt = client.text_generation(
             prompt=prompt_for_llm,
             model=PROMPT_IMPROVER_MODEL,
-            max_new_tokens=128,  # Limit the length of the improved prompt
+            max_new_tokens=128,
             temperature=0.7,
             top_p=0.9,
-            repetition_penalty=1.2, # Encourage diverse output
-            stop_sequences=["</s>"], # stop at end of sentence
+            repetition_penalty=1.2,
+            stop_sequences=["</s>"],
 
         )
 
-        return improved_prompt.strip()  # Remove leading/trailing whitespace
-
+        return improved_prompt.strip()
 
     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
+        print(f"Error improving prompt: {e}")
+        return original_prompt
 
 
 def generate_image(prompt, progress=gr.Progress()):
-    """Generates an image using the InferenceClient and provides progress updates."""
-
     if HF_TOKEN_ERROR:
         raise gr.Error(HF_TOKEN_ERROR)
 
     progress(0, desc="Improving prompt...")
     improved_prompt = improve_prompt(prompt)
 
-    progress(0.2, desc="Sending request to Hugging Face...") # More granular progress
+    progress(0.2, desc="Sending request to Hugging Face...")
     try:
         image = client.text_to_image(improved_prompt, model="black-forest-labs/FLUX.1-schnell")
 
@@ -69,7 +64,7 @@ def generate_image(prompt, progress=gr.Progress()):
         progress(0.8, desc="Processing image...")
         time.sleep(0.5)
         progress(1.0, desc="Done!")
-        return image, improved_prompt  # Return both image and improved prompt
+        return image, 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}"
@@ -77,10 +72,13 @@ def generate_image(prompt, progress=gr.Progress()):
             error_message = f"An error occurred: {e}"
         raise gr.Error(error_message)
 
-
+def pil_to_base64(img):
+    buffered = BytesIO()
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return f"data:image/png;base64,{img_str}"
 
 css = """
-/* ... (Rest of your CSS, unchanged, from the previous response) ... */
 .container {
     max-width: 800px;
     margin: auto;
@@ -94,7 +92,7 @@ css = """
     font-size: 2.5em;
     margin-bottom: 0.5em;
     color: #333;
-    font-family: 'Arial', sans-serif; /* More readable font */
+    font-family: 'Arial', sans-serif;
 }
 .description {
     text-align: center;
@@ -106,15 +104,14 @@ css = """
     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 */
+    display: block;
+    margin: auto;
+    max-width: 100%;
+    height: auto;
+    border-radius: 8px;
+    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
 }
 
-/* Animation for the image appearance - subtle fade-in */
 @keyframes fadeIn {
     from { opacity: 0; transform: translateY(20px); }
     to   { opacity: 1; transform: translateY(0); }
@@ -123,7 +120,6 @@ css = """
     animation: fadeIn 0.8s ease-out;
 }
 
-/* Improve button style */
 .submit-button {
     display: block;
     margin: auto;
@@ -140,7 +136,6 @@ css = """
     background-color: #367c39;
 }
 
-/* Style the error messages */
 .error-message {
     color: red;
     text-align: center;
@@ -148,9 +143,9 @@ css = """
     font-weight: bold;
 }
 label{
-    font-weight: bold; /* Make labels bold */
-    display: block;      /* Each label on its own line */
-    margin-bottom: 0.5em; /* Space between label and input */
+    font-weight: bold;
+    display: block;
+    margin-bottom: 0.5em;
 }
 
 .improved-prompt-display {
@@ -162,6 +157,18 @@ label{
     font-style: italic;
     color: #444;
 }
+.download-link {
+    display: block;
+    text-align: center;
+    margin-top: 10px;
+    color: #4CAF50;
+    text-decoration: none;
+    font-weight: bold;
+}
+
+.download-link:hover{
+  text-decoration: underline;
+}
 """
 
 
@@ -169,7 +176,6 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown(
         """
         # Xylaria Iris Image Generator
-        Enter a text prompt, and we'll enhance it before generating an image!
         """,
         elem_classes="title"
     )
@@ -181,18 +187,22 @@ with gr.Blocks(css=css) as demo:
                 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", show_download_button=False, interactive=False)  # No SVG, not interactive
+                image_output = gr.Image(label="Generated Image", interactive=False)
                 improved_prompt_output = gr.Textbox(label="Improved Prompt", interactive=False, elem_classes="improved-prompt-display")
+                download_link = gr.HTML(visible=False)
 
 
     def on_generate_click(prompt):
         output_group.elem_classes = ["output-section", "animate"]
         image, improved_prompt = generate_image(prompt)
         output_group.elem_classes = ["output-section"]
-        return image, improved_prompt
+        image_b64 = pil_to_base64(image)
+        download_html = f'<a class="download-link" href="{image_b64}" download="generated_image.png">Download Image</a>'
+
+        return image, improved_prompt, download_html
 
-    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])
+    generate_button.click(on_generate_click, inputs=prompt_input, outputs=[image_output, improved_prompt_output, download_link])
+    prompt_input.submit(on_generate_click, inputs=prompt_input, outputs=[image_output, improved_prompt_output, download_link])
 
     gr.Examples(
         [["A dog"],