Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from huggingface_hub import model_info | |
| from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler | |
| def inference(prompt, model, n_images, seed): | |
| # Load the model | |
| info = model_info(model) | |
| model_base = info.cardData["base_model"] | |
| pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float32) | |
| pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe.unet.load_attn_procs(model) | |
| # Load the UI components for progress bar and image grid | |
| progress_bar_ui = st.empty() | |
| with progress_bar_ui.container(): | |
| progress_bar = st.progress(0, text=f"Performing inference on {n_images} images...") | |
| image_grid_ui = st.empty() | |
| # Run inference | |
| result_images = [] | |
| generators = [torch.Generator().manual_seed(i) for i in range(seed, n_images+seed)] | |
| print(f"Inferencing '{prompt}' for {n_images} images.") | |
| for i in range(n_images): | |
| result = pipe(prompt, generator=generators[i], num_inference_steps=25).images[0] | |
| result_images.append(result) | |
| # Start with empty UI elements | |
| progress_bar_ui.empty() | |
| image_grid_ui.empty() | |
| # Update the progress bar | |
| with progress_bar_ui.container(): | |
| value = ((i+1)/(len(dataset))) | |
| progress_bar.progress(value, text=f"{i+1} out of {len(dataset)} images processed.") | |
| # Update the image grid | |
| with image_grid_ui.container(): | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| for i in range(0, len(result_images), 3): | |
| st.image(result_images[i], caption=f"Image - {i+1}") | |
| with col2: | |
| for i in range(1, len(result_images), 3): | |
| st.image(result_images[i], caption=f"Image - {i+2}") | |
| with col3: | |
| for i in range(2, len(result_images), 3): | |
| st.image(result_images[i], caption=f"Image - {i+3}") | |
| def main(): | |
| pass | |
| if __name__ == "__main__": | |
| # --- START UI --- | |
| st.title("Finetune LoRA inference") | |
| with st.form(key='form_parameters'): | |
| prompt = st.text_input("Enter the prompt: ") | |
| model_options = ["asrimanth/person-thumbs-up-plain-lora", "asrimanth/person-thumbs-up-lora", "asrimanth/person-thumbs-up-lora-no-cap"] | |
| current_model = st.selectbox("Choose a model", options=model_options) | |
| col1_inp, col2_inp = st.columns(2) | |
| with col1_inp: | |
| n_images = int(st.number_input("Enter the number of images", min_value=0, max_value=50)) | |
| with col2_inp: | |
| seed_input = int(st.number_input("Enter the seed (default=25)", value=25, min_value=0)) | |
| submitted = st.form_submit_button("Predict") | |
| if submitted: # The form is submitted | |
| inference(prompt, current_model, n_images, seed_input) | |