app.py

import torch
import streamlit as st
from diffusers import AutoPipelineForText2Image
from PIL import Image

# Initialize the model pipeline
@st.cache_resource
def load_pipeline():
    pipeline = AutoPipelineForText2Image.from_pretrained(
        'black-forest-labs/FLUX.1-dev', torch_dtype=torch.float16
    ).to('cuda')
    pipeline.load_lora_weights('hxxxp/bears-nfts', weight_name='lora.safetensors')
    return pipeline

pipeline = load_pipeline()

# Streamlit interface for inputs
st.title("Interactive Image Generation with Streamlit")

# Text prompt input
prompt = st.text_area("Enter your prompt for image generation", "A beautiful landscape with mountains and rivers")

# Image generation parameters
num_images = st.slider("Number of Images", 1, 5, 1)  # User selects how many images to generate
guidance_scale = st.slider("Guidance Scale", 1.0, 10.0, 7.5)  # Control how strict the model should follow the prompt
height = st.slider("Image Height", 256, 1024, 512)  # User can change image height
width = st.slider("Image Width", 256, 1024, 512)  # User can change image width

# Generate image button
if st.button("Generate Image"):
    with st.spinner("Generating..."):
        # Generate images using the model
        images = pipeline(
            prompt, num_inference_steps=50, num_images_per_prompt=num_images, guidance_scale=guidance_scale, height=height, width=width
        ).images

        # Display generated images
        for idx, image in enumerate(images):
            st.image(image, caption=f"Generated Image {idx+1}", use_column_width=True)