Spaces:
Running
Running
File size: 1,697 Bytes
e16aa33 b18dce3 e16aa33 b18dce3 e16aa33 b18dce3 e16aa33 b18dce3 e16aa33 b18dce3 e16aa33 b18dce3 e16aa33 b18dce3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import streamlit as st
import torch
from diffusers import StableDiffusionXLPipeline
from PIL import Image
# Load the model (CPU version)
@st.cache_resource()
def load_model():
pipe = StableDiffusionXLPipeline.from_pretrained(
"cagliostrolab/animagine-xl-4.0",
torch_dtype=torch.float32, # Change to float32 for CPU
use_safetensors=True,
custom_pipeline="lpw_stable_diffusion_xl",
add_watermarker=False
)
return pipe # Do NOT move to CUDA
pipe = load_model()
# Streamlit UI
st.title("Text-to-Image Generator (CPU)")
# User input
prompt = st.text_area("Enter your prompt:", "1girl, arima kana, oshi no ko, looking at viewer, smile")
width = st.slider("Width", 512, 1024, 832, step=64)
height = st.slider("Height", 512, 1280, 1216, step=64)
guidance_scale = st.slider("Guidance Scale", 1.0, 10.0, 6.0, step=0.5)
inference_steps = st.slider("Inference Steps", 10, 50, 25, step=5)
if st.button("Generate Image"):
try:
with st.spinner("Generating image..."):
image = pipe(
prompt,
width=width,
height=height,
guidance_scale=guidance_scale,
num_inference_steps=inference_steps
).images[0]
st.image(image, caption="Generated Image", use_column_width=True)
image_path = "generated_image.png"
image.save(image_path)
with open(image_path, "rb") as file:
st.download_button(label="Download Image", data=file, file_name="generated_image.png", mime="image/png")
except Exception as e:
st.error(f"Error generating image: {e}")
|