Spaces:
Running
Running
import streamlit as st | |
from diffusers import StableDiffusionPipeline | |
import torch | |
from PIL import Image | |
import io | |
import os | |
# Force CPU usage on Spaces | |
os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
def load_model(): | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float32, | |
) | |
pipe.to("cpu") # Force CPU usage | |
return pipe | |
# Streamlit UI | |
st.title("π¨ AI Image Generator (CPU - Hugging Face Spaces)") | |
prompt = st.text_input("Enter your prompt:", | |
"A multi-dimensional futuristic city with glowing lights, fractals, 8K") | |
guidance = st.slider("Creativity (Guidance Scale)", 1.0, 20.0, 7.5) | |
if st.button("Generate"): | |
with st.spinner("Generating image on CPU... please wait β³"): | |
pipe = load_model() | |
image = pipe(prompt, guidance_scale=guidance).images[0] | |
st.image(image, caption="Generated Image", use_column_width=True) | |
buf = io.BytesIO() | |
image.save(buf, format="PNG") | |
st.download_button("Download Image", buf.getvalue(), "generated.png", "image/png") | |