lab10 / app.py
SmallO's picture
Upload 2 files
866879a
# -*- 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.float16)
pipe = pipe.to("cuda")
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()