fix bugs on generate function

app.py

@@ -5,10 +5,12 @@ from PIL import Image
 import gradio as gr
 from datetime import datetime
 
+# Retrieve the Hugging Face token from environment variables
+api_token = os.getenv("HF_TOKEN")
+
 # Debugging: Check if the Hugging Face token is available
-api_token = os.getenv("HF_CTB_TOKEN")
 if not api_token:
-    print("ERROR: Hugging Face token (HF_CTB_TOKEN) is missing. Please set it as an environment variable.")
+    print("ERROR: Hugging Face token (HF_TOKEN) is missing. Please set it as an environment variable.")
 else:
     print("Hugging Face token loaded successfully.")
 
@@ -67,11 +69,14 @@ prompts = [
 def generate_image(prompt_alias, team, model_alias, height, width, num_inference_steps, guidance_scale, seed):
     # Debugging: Check if the token is available
     if not api_token:
-        return None, "ERROR: Hugging Face token (HF_CTB_TOKEN) is missing. Please set it as an environment variable."
+        return None, "ERROR: Hugging Face token (HF_TOKEN) is missing. Please set it as an environment variable."
 
     # Find the selected prompt and model
-    prompt = next(p for p in prompts if p["alias"] == prompt_alias)["text"]
-    model_name = next(m for m in models if m["alias"] == model_alias)["name"]
+    try:
+        prompt = next(p for p in prompts if p["alias"] == prompt_alias)["text"]
+        model_name = next(m for m in models if m["alias"] == model_alias)["name"]
+    except StopIteration:
+        return None, "ERROR: Invalid prompt or model selected."
 
     # Determine the enemy color
     enemy_color = "blue" if team.lower() == "red" else "red"
@@ -87,7 +92,10 @@ def generate_image(prompt_alias, team, model_alias, height, width, num_inference
         seed = random.randint(0, 1000000)
 
     # Initialize the InferenceClient
-    client = InferenceClient(model_name, token=api_token)
+    try:
+        client = InferenceClient(model_name, token=api_token)
+    except Exception as e:
+        return None, f"ERROR: Failed to initialize InferenceClient. Details: {e}"
 
     # Generate the image
     try:
@@ -105,7 +113,10 @@ def generate_image(prompt_alias, team, model_alias, height, width, num_inference
     # Save the image with a timestamped filename
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     output_filename = f"{timestamp}_{model_alias.replace(' ', '_').lower()}_{prompt_alias.replace(' ', '_').lower()}_{team.lower()}.png"
-    image.save(output_filename)
+    try:
+        image.save(output_filename)
+    except Exception as e:
+        return None, f"ERROR: Failed to save image. Details: {e}"
 
     return output_filename, "Image generated successfully!"
 
@@ -132,19 +143,11 @@ with gr.Blocks() as demo:
     # Function to handle button click
     def generate(prompt_alias, team, model_alias, height, width, num_inference_steps, guidance_scale, seed):
         try:
-            # Update status to indicate rendering
-            status_text.update("Rendering image... Please wait.")
-            yield None, "Rendering image... Please wait."  # Yield to update the UI
-
             # Generate the image
             image_path, message = generate_image(prompt_alias, team, model_alias, height, width, num_inference_steps, guidance_scale, seed)
-
-            # Update status and return the image
-            status_text.update(message)
-            yield image_path, message
+            return image_path, message
         except Exception as e:
-            status_text.update(f"An error occurred: {e}")
-            yield None, f"An error occurred: {e}"
+            return None, f"An error occurred: {e}"
 
     # Connect the button to the function
     generate_button.click(