app.py

import os
import gradio as gr
import numpy as np
import random
from pathlib import Path
from PIL import Image
from huggingface_hub import AsyncInferenceClient, InferenceClient
from gradio_client import Client, handle_file
from gradio_imageslider import ImageSlider

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

# Directorio de almacenamiento de imágenes
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)  # Asegura que el directorio exista

def enable_lora(lora_add, basemodel):
    return basemodel if not lora_add else lora_add

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, negative_prompt="",
            seed=42, upscale_factor=upscale_factor, controlnet_scale=0.6,
            controlnet_decay=1, condition_scale=6, tile_width=112,
            tile_height=144, denoise_strength=0.35, num_inference_steps=18,
            solver="DDIM", api_name="/process"
        )
        return result[1] if isinstance(result, list) and len(result) > 1 else None
    except Exception as 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)
    
    image, seed = await generate_image(combined_prompt, model, width, height, scales, steps, seed)

    if isinstance(image, str) and image.startswith("Error"):
        return [image, None]

    image_path = DATA_PATH / f"image_{seed}.jpg"
    image.save(image_path, format="JPEG")

    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")
            return [image_path, DATA_PATH / f"upscale_image_{seed}.jpg"]
        else:
            return [image_path, image_path]
    else:
        return [image_path, image_path]

async def improve_prompt(prompt):
    try:
        instruction = ("With this idea, describe in English a detailed img2vid prompt in a single paragraph of up to 200 characters maximum, developing atmosphere, characters, lighting, and cameras.")
        formatted_prompt = f"{prompt}: {instruction}"
        response = llm_client.text_generation(formatted_prompt, max_new_tokens=200)
        improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()

        return improved_text
    except Exception as e:
        return f"Error mejorando el prompt: {e}"

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 [file["name"] for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"

css = """
#col-container{ margin: 0 auto; max-width: 1024px;}
"""

with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as demo:
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=3):
                output_res = ImageSlider(label="Generadas / Escaladas")
            with gr.Column(scale=2):
                prompt = gr.Textbox(label="Descripción de imagen")
                basemodel_choice = gr.Dropdown(
                    label="Modelo", 
                    choices=["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"],
                    value="black-forest-labs/FLUX.1-schnell"
                )
                lora_model_choice = gr.Dropdown(
                    label="LORA Realismo", 
                    choices=["Shakker-Labs/FLUX.1-dev-LoRA-add-details", "XLabs-AI/flux-RealismLora"],
                    value="XLabs-AI/flux-RealismLora"
                )
                with gr.Row():
                    process_lora = gr.Checkbox(label="Procesar LORA")
                    process_upscale = gr.Checkbox(label="Procesar Escalador")
                improved_prompt = gr.Textbox(label="Prompt Mejorado", interactive=False)
                improve_btn = gr.Button("Mejorar prompt")
                improve_btn.click(fn=improve_prompt, inputs=[prompt], outputs=improved_prompt)
                with gr.Accordion(label="Opciones Avanzadas", open=False):
                    width = gr.Slider(label="Ancho", minimum=512, maximum=1280, step=8, value=1280)
                    height = gr.Slider(label="Alto", minimum=512, maximum=1280, step=8, value=768)
                    upscale_factor = gr.Radio(label="Factor de Escala", choices=[2, 4, 8], value=2)
                    scales = gr.Slider(label="Escalado", minimum=1, maximum=20, step=1, value=10)
                    steps = gr.Slider(label="Pasos", minimum=1, maximum=100, step=1, value=20)
                    seed = gr.Number(label="Semilla", value=-1)
                btn = gr.Button("Generar")
                btn.click(
                    fn=gen, 
                    inputs=[prompt, basemodel_choice, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model_choice, process_lora], 
                    outputs=output_res
                )
        with gr.Row():
            with gr.Column():
                file_list = gr.Gallery(label="Imágenes Guardadas")  # Usar Gallery en lugar de Files
                storage_info = gr.Text(label="Uso de Almacenamiento")
            refresh_btn = gr.Button("Actualizar Galería")
            refresh_btn.click(fn=get_storage, inputs=None, outputs=[file_list, storage_info])
    demo.launch(allowed_paths=[str(DATA_PATH)])