import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

# Load model from Hugging Face
model_id = "harmionestark/GraveImageAI"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Function to generate images
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 GraveImageAI - AI Image Generator")
    with gr.Row():
        prompt = gr.Textbox(label="Enter Prompt", placeholder="A futuristic graveyard under a red sky...")
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")

    generate_button.click(fn=generate_image, inputs=prompt, outputs=output_image)

# Launch app
demo.launch()