import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import numpy as np
import cv2

API_URL = "https://api-inference.huggingface.co/models/dalle-mini/dalle-mini/mega-1-fp16"

def image_to_sketch(image):
    gray = image.convert("L")
    inv = 255 - np.array(gray)
    blur = cv2.GaussianBlur(inv, (21, 21), 0)
    sketch = cv2.divide(np.array(gray), 255 - blur, scale=256)
    return Image.fromarray(sketch)

def generate_sketch(prompt):
    full_prompt = prompt.strip() + ", pencil sketch, line art, black and white"
    response = requests.post(API_URL, json={"inputs": full_prompt})

    try:
        output = response.json()
        if isinstance(output, dict) and output.get("error"):
            return f"API Error: {output['error']}"
        # Newer inference endpoints sometimes return `images` key
        if isinstance(output, list) and "generated_image" in output[0]:
            img_url = output[0]["generated_image"]
            img_data = requests.get(img_url).content
            image = Image.open(BytesIO(img_data))
            return image_to_sketch(image)
        else:
            return "Unexpected response format. Try again later."
    except Exception as e:
        return f"Error: {str(e)}"

gr.Interface(
    fn=generate_sketch,
    inputs=gr.Textbox(placeholder="e.g. a wizard fighting a dragon"),
    outputs="image",
    title="Text to Sketch AI",
    description="Type a description, and get a sketch-style image using DALL·E Mini + OpenCV."
).launch()