Spaces:
Running
Running
File size: 5,881 Bytes
f693d07 aec9996 f693d07 df23192 f693d07 df23192 f693d07 aec9996 18c1250 f693d07 df23192 f693d07 aec9996 f693d07 df23192 18c1250 df23192 18c1250 d2a5152 18c1250 f693d07 df23192 134fe34 f693d07 18c1250 aec9996 df23192 aec9996 df23192 134fe34 df23192 134fe34 f693d07 134fe34 18c1250 f693d07 18c1250 df23192 18c1250 aec9996 d99c194 7c6a22d aec9996 ab442ba 434794f ab442ba aec9996 ab442ba df23192 aec9996 df23192 f693d07 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import os
import zipfile
import numpy as np
import random
from pathlib import Path
from PIL import Image
import streamlit as st
from huggingface_hub import AsyncInferenceClient
import asyncio
MAX_SEED = np.iinfo(np.int32).max
client = AsyncInferenceClient()
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)
PREDEFINED_SEED = random.randint(0, MAX_SEED)
async def generate_image(prompt, width, height, seed, model_name):
try:
if seed == -1:
seed = PREDEFINED_SEED
seed = int(seed)
image = await client.text_to_image(
prompt=prompt, height=height, width=width, model=model_name
)
return image, seed
except Exception as e:
return f"Error al generar imagen: {e}", 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, width, height, model_name):
combined_prompt = prompt
seed = PREDEFINED_SEED
progress_bar = st.progress(0)
image, seed = await generate_image(combined_prompt, width, height, seed, model_name)
progress_bar.progress(100)
if isinstance(image, str) and image.startswith("Error"):
progress_bar.empty()
return [image, None]
image_path = save_image(image, seed)
prompt_file_path = save_prompt(combined_prompt, seed)
return [str(image_path), str(prompt_file_path)]
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 = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
usage = sum([file.stat().st_size for file in files])
return [str(file.resolve()) for file 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_all_images():
try:
files = [file for file in DATA_PATH.glob("*.jpg")]
prompts = [file for file in DATA_PATH.glob("*.txt")]
for file in files + prompts:
os.remove(file)
st.success("Todas las imágenes y prompts han sido borrados.")
except Exception as e:
st.error(f"Error al borrar archivos: {e}")
def download_images_as_zip():
zip_path = DATA_PATH / "images.zip"
with zipfile.ZipFile(zip_path, 'w') as zipf:
for file in DATA_PATH.glob("*.jpg"):
zipf.write(file, arcname=file.name)
with open(zip_path, "rb") as zip_file:
st.download_button(label="Descargar imágenes en .zip", data=zip_file, file_name="images.zip", mime="application/zip")
def main():
st.set_page_config(layout="wide")
st.title("Generador de Imágenes")
st.warning("Este espacio contiene contenido que no es adecuado para todas las audiencias. Se recomienda discreción.")
agree = st.checkbox("Soy mayor de 18 años y entiendo que el contenido puede no ser apropiado.")
if agree:
prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
model_option = st.sidebar.selectbox("Modelo", ["enhanceaiteam/Flux-uncensored", "enhanceaiteam/Flux-Uncensored-V2"])
width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
if st.sidebar.button("Generar Imagen"):
with st.spinner("Generando imagen..."):
result = asyncio.run(gen(prompt, width, height, model_option))
image_paths = result[0]
prompt_file = result[1]
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 st.button("Borrar todas las imágenes"):
delete_all_images()
download_images_as_zip()
else:
st.error("No puedes acceder a este espacio hasta que confirmes que eres mayor de edad.")
if __name__ == "__main__":
main() |