flux4 / app.py
salomonsky's picture
Update app.py
8c45d22 verified
raw
history blame
7.98 kB
import os
import numpy as np
import random
from pathlib import Path
from PIL import Image
import streamlit as st
from huggingface_hub import InferenceClient, AsyncInferenceClient
from gradio_client import Client, handle_file
import asyncio
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
client = AsyncInferenceClient()
llm_client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)
def enable_lora(lora_add, basemodel):
return lora_add if lora_add else basemodel
async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
try:
if seed == -1:
seed = random.randint(0, MAX_SEED)
seed = int(seed)
image = await client.text_to_image(
prompt=combined_prompt, height=height, width=width, guidance_scale=scales,
num_inference_steps=steps, model=model
)
return image, seed
except Exception as e:
return f"Error al generar imagen: {e}", None
def get_upscale_finegrain(prompt, img_path, upscale_factor):
try:
client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
result = client.predict(
input_image=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor
)
return result[1] if isinstance(result, list) and len(result) > 1 else None
except Exception as e:
return None
def save_prompt(prompt_text, seed):
try:
prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
with open(prompt_file_path, "w") as prompt_file:
prompt_file.write(prompt_text)
return prompt_file_path
except Exception as e:
st.error(f"Error al guardar el prompt: {e}")
return None
async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora):
model = enable_lora(lora_model, basemodel) if process_lora else basemodel
improved_prompt = await improve_prompt(prompt)
combined_prompt = f"{prompt} {improved_prompt}"
if seed == -1:
seed = random.randint(0, MAX_SEED)
seed = int(seed)
progress_bar = st.progress(0)
image, seed = await generate_image(combined_prompt, model, width, height, scales, steps, seed)
progress_bar.progress(50)
if isinstance(image, str) and image.startswith("Error"):
progress_bar.empty()
return [image, None, combined_prompt]
image_path = save_image(image, seed)
prompt_file_path = save_prompt(combined_prompt, seed)
if process_upscale:
upscale_image_path = get_upscale_finegrain(combined_prompt, image_path, upscale_factor)
if upscale_image_path:
upscale_image = Image.open(upscale_image_path)
upscale_image.save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
progress_bar.progress(100)
image_path.unlink()
return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
else:
progress_bar.empty()
return [str(image_path), str(prompt_file_path)]
else:
progress_bar.progress(100)
return [str(image_path), str(prompt_file_path)]
async def improve_prompt(prompt):
try:
instruction = ("With this idea, describe in English a detailed txt2img prompt in 300 characters at most...")
formatted_prompt = f"{instruction}"
response = llm_client.text_generation(formatted_prompt, max_new_tokens=300)
improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()
return improved_text[:300] if len(improved_text) > 300 else improved_text
except Exception as e:
return f"Error mejorando el prompt: {e}"
def save_image(image, seed):
try:
image_path = DATA_PATH / f"image_{seed}.jpg"
image.save(image_path, format="JPEG")
return image_path
except Exception as e:
st.error(f"Error al guardar la imagen: {e}")
return None
def get_storage():
files = [{"name": str(file.resolve()), "size": file.stat().st_size} for file in DATA_PATH.glob("*.jpg") if file.is_file()]
usage = sum([f['size'] for f in files])
return [f["name"] for f in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
def get_prompts():
prompt_files = [file for file in DATA_PATH.glob("*.txt") if file.is_file()]
return {file.stem.replace("prompt_", ""): file for file in prompt_files}
def delete_image(image_path):
try:
if Path(image_path).exists():
Path(image_path).unlink()
st.success(f"Imagen {image_path} borrada.")
else:
st.error("El archivo de imagen no existe.")
except Exception as e:
st.error(f"Error al borrar la imagen: {e}")
def main():
st.set_page_config(layout="wide")
st.title("Generador de Imágenes FLUX")
prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=200)
with st.sidebar.expander("Opciones avanzadas", expanded=False):
basemodel = st.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"])
lora_model = st.selectbox("LORA Realismo", ["Shakker-Labs/FLUX.1-dev-LoRA-add-details", "XLabs-AI/flux-RealismLora"])
format_option = st.selectbox("Formato", ["9:16", "16:9"])
process_lora = st.checkbox("Procesar LORA")
process_upscale = st.checkbox("Procesar Escalador")
upscale_factor = st.selectbox("Factor de Escala", [2, 4, 8], index=0)
scales = st.slider("Escalado", 1, 20, 10)
steps = st.slider("Pasos", 1, 100, 20)
seed = st.number_input("Semilla", value=-1)
if format_option == "9:16":
width = 720
height = 1280
else:
width = 1280
height = 720
if st.sidebar.button("Generar Imagen"):
with st.spinner("Mejorando y generando imagen..."):
improved_prompt = asyncio.run(improve_prompt(prompt))
st.session_state.improved_prompt = improved_prompt
prompt_to_use = st.session_state.get('improved_prompt', prompt)
result = asyncio.run(gen(prompt_to_use, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora))
image_paths = result[0]
prompt_file = result[1]
st.write(f"Image paths: {image_paths}")
if image_paths:
if Path(image_paths).exists():
st.image(image_paths, caption="Imagen Generada")
else:
st.error("El archivo de imagen no existe.")
if prompt_file and Path(prompt_file).exists():
prompt_text = Path(prompt_file).read_text()
st.write(f"Prompt utilizado: {prompt_text}")
else:
st.write("El archivo del prompt no está disponible.")
files, usage = get_storage()
st.text(usage)
cols = st.columns(6)
prompts = get_prompts()
for idx, file in enumerate(files):
with cols[idx % 6]:
image = Image.open(file)
prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
st.image(image, caption=f"Imagen {idx+1}")
st.write(f"Prompt: {prompt_text}")
if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
try:
os.remove(file)
if prompt_file:
os.remove(prompt_file)
st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
except Exception as e:
st.error(f"Error al borrar la imagen o prompt: {e}")
if __name__ == "__main__":
main()