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()