Update app.py

app.py

@@ -1,44 +1,48 @@
 import gradio as gr
 import requests
-from io import BytesIO
+import os
 from PIL import Image
 import os
-# Load Stable Diffusion model
-#HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None
-#model = gr.load("models/stabilityai/stable-diffusion-3.5-large", hf_token=HF_TOKEN)
-model = gr.Interface.load("models/stabilityai/stable-diffusion-3.5-large")
-print(model)
-# Chevereto API details
-CHEVERETO_API_URL = os.environ.get("api_url")  # Replace with your Chevereto API endpoint
-CHEVERETO_API_KEY =  os.environ.get("API_KEY")  # Replace with your API key
-CHEVERETO_ALBUM_ID =  os.environ.get("ALBUM_ID")  # Replace with the album ID where you want to store images
-
-# Upload image to Chevereto
-def upload_to_chevereto(image: Image.Image) -> str:
-    # Convert image to bytes
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    files = {"source": buffered.getvalue()}  # Convert to binary for upload
-    
-    # Set API parameters
-    data = {
-        "key": CHEVERETO_API_KEY,
-        "format": "json",
-        "album": CHEVERETO_ALBUM_ID  # Specify album ID
-    }
+# Load environment variables
+CHEVERETO_API_URL = os.environ.get("API_URL")  # Your Chevereto API endpoint
+CHEVERETO_API_KEY = os.environ.get("API_KEY")  # Your Chevereto API Key
+CHEVERETO_ALBUM_ID = os.environ.get("ALBUM_ID")  # Your Album ID
 
-    # Upload request
-    response = requests.post(CHEVERETO_API_URL, files=files, data=data)
+# Load the Gradio model
+model = gr.Interface.load("models/stabilityai/stable-diffusion-3.5-large")
 
+def upload_to_chevereto(img_path: str) -> str:
+    """Uploads an image to Chevereto and returns the public image URL."""
+    with open(img_path, "rb") as image_file:
+        files = {"source": image_file}
+        data = {
+            "key": CHEVERETO_API_KEY,
+            "format": "json",
+            "album": CHEVERETO_ALBUM_ID  # Include the album ID
+        }
+        response = requests.post(CHEVERETO_API_URL, files=files, data=data)
+    
     if response.status_code == 200:
-        return response.json()["image"]["url"]  # Return public URL of uploaded image
-    return "Error uploading image"
+        return response.json().get("image", {}).get("url")  # Return the image URL
+    else:
+        return f"Error uploading image: {response.text}"  # Debugging error message
 
-# Generate and upload the image
 def generate_image_and_upload(prompt: str) -> str:
-    #img = model.submit([prompt])  # Get generated image
-    img = model.predict(prompt)
-    return upload_to_chevereto(img)  # Upload and return the URL
+    """Generates an image using the AI model and uploads it to Chevereto."""
+    img = model(prompt)  # Generate image
+
+    if isinstance(img, list):  
+        img = img[0]  # If Gradio returns a list, get the first image
+
+    if isinstance(img, Image.Image):
+        img_path = "generated_image.png"
+        img.save(img_path, format="PNG")  # Save locally before uploading
+
+        return upload_to_chevereto(img_path)  # Upload and return URL
+    else:
+        return "Error: Model did not return a valid image"
+
+
 
 # Define Gradio API returning the Chevereto public URL
 iface = gr.Interface(fn=generate_image_and_upload, inputs="text", outputs="text")