app.py

import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random
import os
import json
import io
import uuid
from gradio_client import Client as client_gradio
from supabase import create_client, Client
from datetime import datetime

# Inicializa Supabase
url: str = os.getenv('SUPABASE_URL')
key: str = os.getenv('SUPABASE_KEY')
supabase: Client = create_client(url, key)

# Obtém token da Hugging Face
hf_token = os.getenv("HF_TOKEN")

# Inicializa o modelo base FLUX.1-dev
base_model = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(
    base_model, 
    torch_dtype=torch.float16,  # Pode testar torch.float32 se a GPU permitir
    use_safetensors=True
)

# Move o modelo para GPU
pipe.to("cuda")

# Definição dos LoRA e Trigger Words
lora_models = {
    "studio": {
        "repo": "vcollos/Paula2",
        "weights": "Paula P.safetensors",
        "trigger_word": ""  # Sem trigger word específica
    },
    "Nanda": {
        "repo": "vcollos/Nanda",
        "weights": "lora.safetensors",
        "trigger_word": ""
    }
}

# Carrega os LoRAs disponíveis
for name, details in lora_models.items():
    try:
        pipe.load_lora_weights(details["repo"], weight_name=details["weights"], adapter_name=name)
        print(f"✅ LoRA {name} carregado")
    except Exception as e:
        print(f"❌ Erro ao carregar o LoRA {name}: {e}")

# Define seed máximo
MAX_SEED = 2**32 - 1

def upload_image_to_supabase(image, filename):
    """ Faz upload da imagem para o Supabase Storage e retorna a URL pública. """
    img_bytes = io.BytesIO()
    image.save(img_bytes, format="PNG")
    img_bytes.seek(0)  # Move para o início do arquivo

    storage_path = f"images/{filename}"

    try:
        # Faz upload da imagem para o Supabase
        supabase.storage.from_("images").upload(storage_path, img_bytes.getvalue(), {"content-type": "image/png"})

        # Retorna a URL pública da imagem
        base_url = f"{url}/storage/v1/object/public/images"
        return f"{base_url}/{filename}"
    except Exception as e:
        print(f"❌ Erro no upload da imagem: {e}")
        return None

@spaces.GPU(duration=80)
def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # Trunca o prompt para 77 tokens para evitar erro do CLIP
    prompt_tokens = prompt.split()[:77]
    prompt = " ".join(prompt_tokens)

    # Define qual LoRA usar com base na seleção do usuário
    selected_loras = []
    adapter_weights = []

    if lora_option == "studio":
        selected_loras.append("studio")
        adapter_weights.append(lora_scale_1)
    elif lora_option == "Nanda":
        selected_loras.append("Nanda")
        adapter_weights.append(lora_scale_2)
    elif lora_option == "Ambos":
        selected_loras = ["studio", "Nanda"]
        adapter_weights = [lora_scale_1, lora_scale_2]

    pipe.set_adapters(selected_loras, adapter_weights)

    # Adiciona trigger words apenas se Nanda estiver ativado
    if "Nanda" in selected_loras:
        prompt = f"{lora_models['Nanda']['trigger_word']} {prompt}"

    # Gera a imagem com precisão de 16 bits para tentar melhorar a nitidez
    with torch.autocast("cuda"):
        image = pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator
        ).images[0]

    # Define um nome único para a imagem
    filename = f"image_{seed}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.png"

    try:
        image_url = upload_image_to_supabase(image, filename)
        if image_url:
            print(f"✅ Imagem salva no Supabase: {image_url}")
        else:
            print("❌ Erro: URL da imagem retornou None")
            return image, seed
    except Exception as e:
        print(f"❌ Erro ao fazer upload da imagem: {e}")
        return image, seed

    # Salva todos os metadados no Supabase
    try:
        response = supabase.table("images").insert({
            "prompt": prompt,
            "cfg_scale": cfg_scale,
            "steps": steps,
            "seed": seed,
            "lora_option": lora_option,
            "lora_scale_1": lora_scale_1,
            "lora_scale_2": lora_scale_2,
            "image_url": image_url,
            "created_at": datetime.utcnow().isoformat()
        }).execute()
        
        if response.data:
            print("✅ Metadados salvos no Supabase")
        else:
            print("❌ Erro: Resposta vazia do Supabase")
        
    except Exception as e:
        print(f"❌ Erro ao salvar metadados no Supabase: {e}")

    return image, seed

# Interface Gradio
gr_theme = os.getenv("THEME")
with gr.Blocks(theme=gr_theme) as app:
    gr.Markdown("# studio Image Generator")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.TextArea(label="Prompt", placeholder="Digite um prompt (máx 77 caracteres)", lines=3)
            generate_button = gr.Button("Gerar")
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
            steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=32)
            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=768)  
            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=1024)  
            randomize_seed = gr.Checkbox(False, label="Randomize seed")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=556215326)
            lora_option = gr.Radio(["Nenhum", "studio", "Nanda", "Ambos"], label="Escolha o LoRA", value="Ambos")
            lora_scale_1 = gr.Slider(label="LoRA Scale (studio)", minimum=0, maximum=1, step=0.01, value=0.1)
            lora_scale_2 = gr.Slider(label="LoRA Scale (Nanda)", minimum=0, maximum=1, step=0.01, value=1)

        with gr.Column(scale=2):
            result = gr.Image(label="Generated Image")

    generate_button.click(
        run_lora,
        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2],
        outputs=[result, seed],
    )

app.queue()
app.launch(share=True)