File size: 3,763 Bytes
cab60db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import streamlit as st
from diffusers import DiffusionPipeline
import torch
from PIL import Image
import io
import gc # Import garbage collection
# Заголовок приложения
st.title("Генератор изображений с LCM Dreamshaper")
st.write("Используйте эту модель для быстрой генерации изображений на CPU")
# Создаем область для настроек
with st.sidebar:
st.header("Настройки")
prompt = st.text_area("Введите ваш запрос:", "hoaxx kitty", height=100)
num_inference_steps = st.slider(
"Количество шагов инференса:",
min_value=1,
max_value=50,
value=5,
help="Больше шагов = выше качество, но медленнее"
)
guidance_scale = st.slider(
"Guidance Scale:",
min_value=1.0,
max_value=15.0,
value=8.0,
step=0.5,
help="Насколько строго модель следует промпту"
)
lcm_origin_steps = st.slider(
"LCM Origin Steps:",
min_value=1,
max_value=50,
value=35
)
generate_button = st.button("Сгенерировать изображение")
# Загружаем модель при первом запуске
@st.cache_resource
def load_model():
pipe = DiffusionPipeline.from_pretrained(
"SimianLuo/LCM_Dreamshaper_v7",
torch_dtype=torch.float32
)
pipe.to("cpu")
pipe.enable_attention_slicing()
pipe.safety_checker = None
return pipe
# Функция для генерации изображения
def generate_image(pipe, prompt, steps, guidance, lcm_steps):
try:
with torch.inference_mode():
images = pipe(
prompt=prompt,
num_inference_steps=steps,
guidance_scale=guidance,
lcm_origin_steps=lcm_steps,
output_type="pil"
).images
return images[0]
except Exception as e:
st.error(f"Error generating image: {e}")
return None
# Загружаем модель
pipe = load_model()
# Отображаем прогресс и результат
if generate_button:
with st.spinner("Генерация изображения..."):
# Создаем место для вывода изображения
result_container = st.empty()
# Генерируем изображение
image = generate_image(
pipe,
prompt,
num_inference_steps,
guidance_scale,
lcm_origin_steps
)
# Показываем результат
if image: # Only display if image generation was successful
result_container.image(image, caption=f"Результат для: {prompt}", use_container_width=True) # Use use_container_width
# Предлагаем скачать
buf = io.BytesIO()
image.save(buf, format="PNG")
byte_im = buf.getvalue()
st.download_button(
label="Скачать изображение",
data=byte_im,
file_name="generated_image.png",
mime="image/png"
)
gc.collect() # Garbage collection after image generation
# Инструкции по использованию
if not generate_button:
st.write("👈 Настройте параметры в боковой панели и нажмите 'Сгенерировать изображение'") |