Update app.py

app.py

@@ -4,6 +4,8 @@ from dotenv import load_dotenv
 from huggingface_hub import InferenceApi
 from PIL import Image
 from io import BytesIO
+import requests
+import base64
 
 # Load environment variables from the .env file
 load_dotenv()
@@ -28,9 +30,25 @@ if st.button("Generate Image"):
             # Make request to the Hugging Face model
             output = inference(inputs=prompt)
 
-            # Convert the output to an image
-            image = Image.open(BytesIO(output))
-            st.image(image, caption="Generated Image", use_column_width=True)
+            # Check if the output is a valid PIL image (already in image format)
+            if isinstance(output, Image.Image):
+                image = output
+            # Check if the output contains a base64-encoded string
+            elif isinstance(output, dict) and 'generated_image_base64' in output:
+                # Decode the base64 string to bytes
+                image_data = base64.b64decode(output['generated_image_base64'])
+                image = Image.open(BytesIO(image_data))
+            # If output contains an image URL
+            elif isinstance(output, dict) and 'generated_image_url' in output:
+                response = requests.get(output['generated_image_url'])
+                image = Image.open(BytesIO(response.content))
+            else:
+                st.error("Unexpected output format from the inference API.")
+                image = None
+
+            # Display the image
+            if image:
+                st.image(image, caption="Generated Image", use_column_width=True)
         except Exception as e:
             st.error(f"Error: {str(e)}")
     else: