from diffusers import StableDiffusionPipeline
import torch
import multiprocessing


# Função para dividir um texto longo em partes menores
def dividir_prompt(texto, limite_tokens=77):
    palavras = texto.split()
    partes = []
    while palavras:
        parte = []
        while palavras and len(" ".join(parte + [palavras[0]]).split()) <= limite_tokens:
            parte.append(palavras.pop(0))
        partes.append(" ".join(parte))
    return partes


# Função que vai processar um lote de imagens em um dispositivo específico (GPU ou CPU)
def processar_lote(lote_prompts, dispositivo, caminho_saida_prefixo="imagem_gerada"):
    print(f"Processando no dispositivo {dispositivo}...")
    modelo = "runwayml/stable-diffusion-v1-5"
    pipeline = StableDiffusionPipeline.from_pretrained(modelo, torch_dtype=torch.float32)

    # Enviar o pipeline para o dispositivo adequado (GPU ou CPU)
    pipeline.to(dispositivo)

    for i, prompt in enumerate(lote_prompts):
        # Gerar a imagem para cada prompt
        with torch.no_grad():  # Desativa o cálculo de gradientes para otimizar a memória
            imagem = pipeline(prompt).images[0]

        # Salvar a imagem gerada
        caminho_saida = f"{caminho_saida_prefixo}_{dispositivo}_{i}.png"
        imagem.save(caminho_saida)
        print(f"Imagem salva em: {caminho_saida}")


# Função para dividir os prompts entre GPU e CPU e processar em paralelo
def gerar_imagens_em_paralelo(prompts, num_lotes=2):
    if len(prompts) < num_lotes:
        num_lotes = len(prompts)  # Não criar mais lotes do que o número de prompts

    # Dividir os prompts em lotes
    lotes = [prompts[i:i + len(prompts) // num_lotes] for i in range(0, len(prompts), len(prompts) // num_lotes)]

    # Garantir que o último lote tenha o resto dos prompts (caso haja uma divisão desigual)
    if len(lotes) > num_lotes:
        lotes[num_lotes - 1].extend(lotes[-1])
        del lotes[-1]

    # Criar uma lista de dispositivos, começando com a GPU
    dispositivos = ["cuda", "cpu"]

    # Atribuir dispositivos aos lotes
    dispositivos_para_lotes = []
    for i in range(len(lotes)):
        dispositivo = dispositivos[i % len(dispositivos)]  # Alternar entre "cuda" e "cpu"
        dispositivos_para_lotes.append(dispositivo)

    # Agora, distribuímos os lotes de forma balanceada entre a GPU e a CPU
    processos = []
    for i, lote in enumerate(lotes):
        dispositivo = dispositivos_para_lotes[i]  # Usa o dispositivo atribuído ao lote
        p = multiprocessing.Process(target=processar_lote, args=(lote, dispositivo))
        processos.append(p)
        p.start()

    # Aguardar todos os processos terminarem
    for p in processos:
        p.join()

if __name__ == '__main__':
    # Exemplo de uso com o texto sobre o gado Nelore
    prompt_original = (
        "A White Nelore Cattle, In The Green Pasture."
    )

    # Dividir o prompt em partes menores para que caibam no modelo
    prompts_divididos = dividir_prompt(prompt_original)

    gerar_imagens_em_paralelo(prompts_divididos)