import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the model (will take time and ~6GB+ RAM)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    use_auth_token=True  # Optional if you already logged in with huggingface-cli
).to("cpu")

# Inference function
def generate_image(prompt):
    image = pipe(prompt, num_inference_steps=15).images[0]
    return image

# Gradio app
gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(placeholder="A dragon playing guitar under a disco ball"),
    outputs="image",
    title="CJ's Image Generator",
    description="Enter a prompt and summon an image like an AI sorcerer with CPU lag."
).launch(share=True)