| import streamlit as st |
| from diffusers import StableDiffusionPipeline |
| import torch |
|
|
| @st.cache_resource(show_spinner=False) |
| def load_model(): |
| model_id = "CompVis/stable-diffusion-v1-4" |
| pipe = StableDiffusionPipeline.from_pretrained( |
| model_id, revision="fp16", torch_dtype=torch.float16 |
| ) |
| pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu") |
| return pipe |
|
|
| def main(): |
| st.title("🖌️ Stickman Sticker Generator") |
| st.write("Enter a description and get a stickman-style image!") |
|
|
| prompt = st.text_input("Enter your prompt", "simple stickman drawing doing yoga, line art") |
|
|
| if st.button("Generate"): |
| with st.spinner("Generating..."): |
| pipe = load_model() |
| image = pipe(prompt).images[0] |
| st.image(image, caption="Generated Stickman Sticker", use_column_width=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
