lab10 / app.py
SmallO's picture
Update app.py
14fb55e
raw
history blame
1.14 kB
# -*- coding: utf-8 -*-
"""Untitled3.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1eg0M7FRTlmONVU_5IaZH-EKADvdtODmT
"""
#!pip install torch
#!pip install diffusers transformers accelerate torch
#!pip install streamlit
from diffusers import StableDiffusionPipeline
import torch
import streamlit as st
def pic_mo(prom):
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32) # Change to float32
# No need to specify .to("cuda") if running on CPU
prompt = str(prom)
image = pipe(prompt).images[0]
return image
def main():
st.title("Text to Image with runwayml/stable-diffusion-v1-5")
# Input text prompt
title = st.text_input('Write a prompt', "")
if st.button('Generate Image'):
# Call the pic_mo function to generate an image
generated_image = pic_mo(title)
# Display the generated image
st.image(generated_image, caption='Generated Image', use_column_width=True)
if __name__ == '__main__':
main()