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")