Update app.py

app.py

@@ -1,32 +1,29 @@
 import os
+import time
 import requests
 import gradio as gr
 import uuid
-import os
 
 # Hugging Face API
-HF_API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
 HF_API_KEY = os.getenv("HF_TOKEN")
 
 # Chevereto API
 CHEVERETO_API_URL = os.getenv("API_URL")
 CHEVERETO_API_KEY = os.getenv("API_KEY")
 CHEVERETO_ALBUM_ID = os.getenv("ALBUM_ID")
-def generate_random_filename(extension: str) -> str:
-    """Generates a random filename with the given extension"""
-    return str(uuid.uuid4()) + extension  # Example: 'b3d50c6d-bfdb-4c4b-b8c4-3a2e3c3c07d9.png'
-# Set your app's secret API key here (you can also load it from the environment)
-EXPECTED_API_KEY = os.getenv("API_PASS")  # or use a hardcoded key for now
+EXPECTED_API_KEY = os.getenv("API_PASS")
 
 HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
 
+def generate_random_filename(extension: str) -> str:
+    return str(uuid.uuid4()) + extension
+
 def validate_api_key(api_key: str) -> bool:
-    """Validates the provided API key"""
     return api_key == EXPECTED_API_KEY
 
-def generate_image(prompt: str):
-    """Generates an image using Hugging Face's API"""
-    response = requests.post(HF_API_URL, headers=HEADERS, json={"inputs": prompt})
+def generate_image(prompt: str, model_name: str):
+    hf_api_url = f"https://api-inference.huggingface.co/models/{model_name}"
+    response = requests.post(hf_api_url, headers=HEADERS, json={"inputs": prompt})
     
     if response.status_code == 200:
         image_data = response.content
@@ -36,15 +33,15 @@ def generate_image(prompt: str):
         with open(image_path, "wb") as f:
             f.write(image_data)
 
+        time.sleep(10)  # 10-sekundowa przerwa
         return image_path
     else:
         return f"Error generating image: {response.text}"
 
 def upload_to_chevereto(image_path: str):
-    """Uploads the generated image to Chevereto and returns the URL"""
     with open(image_path, "rb") as img_file:
         files = {"source": img_file}
-        data = {"key": CHEVERETO_API_KEY, "format": "json", "album_id": CHEVERETO_ALBUM_ID,"expiration": "PT10M"}
+        data = {"key": CHEVERETO_API_KEY, "format": "json", "album_id": CHEVERETO_ALBUM_ID, "expiration": "PT10M"}
         
         response = requests.post(CHEVERETO_API_URL, files=files, data=data)
     
@@ -53,23 +50,23 @@ def upload_to_chevereto(image_path: str):
     else:
         return f"Error uploading image: {response.text}"
 
-def generate_and_upload(prompt: str, api_key: str):
-    """Validates API key, generates image, and uploads it to Chevereto"""
+def generate_and_upload(prompt: str, model_name: str, api_key: str):
     if not validate_api_key(api_key):
-        return "Invalid or missing API key", ""  # Return an error message
+        return "Invalid or missing API key", ""
     
-    img_path = generate_image(prompt)
+    img_path = generate_image(prompt, model_name)
     if "Error" in img_path:
-        return img_path, ""  # Return error message and empty string for URL
+        return img_path, ""
     
     url = upload_to_chevereto(img_path)
-    return img_path, url  # Return the image path and the Chevereto URL
+    return img_path, url
 
 # Gradio UI
 with gr.Blocks() as app:
-    gr.Markdown("<h1 style='text-align: center; color: red;'>🚀 Star Trek AI Image Generator</h1>")
+    gr.Markdown("<h1 style='text-align: center; color: red;'>🚀 SD API on Tasia Space 🚀</h1>")
     with gr.Row():
         prompt_input = gr.Textbox(label="Enter a prompt:", placeholder="A futuristic spaceship...")
+        model_input = gr.Textbox(label="Enter model name:", placeholder="stabilityai/stable-diffusion-3.5-large")
         api_key_input = gr.Textbox(label="Enter your API key:", type="password", placeholder="Your secret API key")
         generate_button = gr.Button("Generate & Upload")
     
@@ -78,7 +75,7 @@ with gr.Blocks() as app:
 
     generate_button.click(
         fn=generate_and_upload,
-        inputs=[prompt_input, api_key_input],
+        inputs=[prompt_input, model_input, api_key_input],
         outputs=[image_output, url_output]
     )