import gradio as gr
import requests
from io import BytesIO
import base64
import os

HF_TOKEN = os.getenv("HF_TOKEN")  # keep token secret

def generate_image(prompt, seed=42):
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}"
    }
    payload = {
        "inputs": prompt,
        "options": {"wait_for_model": True}
    }
    response = requests.post(
        "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4",
        headers=headers,
        json=payload
    )
    image = BytesIO(response.content)
    encoded = base64.b64encode(image.getvalue()).decode()
    return f"data:image/png;base64,{encoded}"

# Use `interface.launch(share=True)` to expose this endpoint to your JS app
app = gr.Interface(fn=generate_image, inputs=["text", "number"], outputs="image")
app.launch()