image_creator / app.py
KKD1998's picture
Update app.py
f265ab6 verified
raw
history blame contribute delete
742 Bytes
import torch
import streamlit as st
from langchain.llms import HuggingFaceHub
from diffusers import FluxPipeline
st.set_page_config(page_title="Image Creator")
st.title("Image Creator Pro...🤖")
# Create a text input box
user_input = st.text_input("Enter some text:")
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
prompt = user_input
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")