from diffusers import FluxPipeline
import torch


BASE_PATH = "/Users/danicosta/Desktop/Flux2/"


print("Cargando pipeline completo desde HuggingFace...")


# 1. Cargar el pipeline con caché local (sin descargar de nuevo)
pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.float32,
    device="mps",
    cache_dir=BASE_PATH + "cache"  # ← Guarda la caché en tu carpeta local
)
# Habilitar offload de CPU para liberar VRAM
## pipeline.enable_model_cpu_offload()

print("Cargando LoRA local...")

# 2. Cargar el LoRA Uncensored desde tu carpeta local
pipeline.load_lora_weights(
    BASE_PATH + "lora.safetensors", 
    weight_name="lora.safetensors"
)


print("Modelos cargados correctamente.")

# 3. Generar imagen con menos pasos (más rápido)
image = pipeline(
    "a hyper realistic picture",
    num_inference_steps=20,  # ← CAMBIO: de 50 a 28 pasos (ahorra ~40% del tiempo)
).images[0]
image.show()

print("Imagen generada y mostrada en la ventana predeterminada de macOS.")