app.py

import requests
from io import BytesIO
from PIL import Image
import gradio as gr
import os

API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
api_key = os.environ.get('API_KEY', 'YOUR_API_KEY_HERE')
headers = {"Authorization": f"Bearer {api_key}"}

class QueryError(Exception):
    pass

def query(payload):
    try:
        assert isinstance(payload, dict)
        response = requests.post(API_URL, headers=headers, json=payload)
        
        if not str(response.status_code).startswith("2"):
            raise QueryError(f"Query failed! Response status code was '{response.status_code}'")
        
        return response.content
        
    except AssertionError:
        print("Invalid Payload Error: Please provide a dictionary.")
    except requests.exceptions.RequestException as e:
        print("Request Failed: ", e)
    except QueryError as qe:
        print(qe)
    except Exception as ex:
        print("Unknown Error occurred: ", ex)

def generate_images_from_prompt(prompt_text, num_images):
    images = []
    for _ in range(num_images):
        image_bytes = query({"inputs": prompt_text})
        img = Image.open(BytesIO(image_bytes))
        images.append(img)
    return images