Update app.py

app.py

@@ -18,7 +18,6 @@ if not HUGGINGFACE_TOKEN:
     raise ValueError("Hugging Face API token not found. Please set it in the environment variables.")
 
 headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}
-print(HUGGINGFACE_TOKEN)
 
 # Function to query Hugging Face model API
 def query(payload):
@@ -34,7 +33,7 @@ def query(payload):
 def generate_image():
     try:
         # Send a request to the Hugging Face model
-        image_bytes = query({"inputs": "cat with dog"})
+        image_bytes = query({"inputs": "dog with cat"})
         
         # Open the image
         image = Image.open(io.BytesIO(image_bytes))
@@ -49,9 +48,18 @@ def generate_image():
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Define a root route to handle requests to "/"
-@app.get("/")
-def read_root(logs: str = None):
-    if logs == "container":
-        return {"logs": "Container logs not available."}
-    return {"message": "Welcome to the API! Use /generate-image to generate images."}
+# Event handler to generate the image on startup
+@app.on_event("startup")
+async def startup_event():
+    try:
+        # Automatically call the generate image logic on startup
+        print("Generating image on startup...")
+        image_bytes = query({"inputs": "dog with cat"})
+        
+        # Save the image locally as 'generated_image.png' for verification
+        with open("generated_image.png", "wb") as f:
+            f.write(image_bytes)
+        print("Image generated and saved as 'generated_image.png'.")
+
+    except Exception as e:
+        print(f"Error during image generation: {e}")