# history

source:
  import sys, os
  FN_ROOT = os.environ.get("FN_REGISTRY_ROOT", "/home/lucas/fn_registry")
  sys.path.insert(0, os.path.join(FN_ROOT, "python/functions/ml"))

  from cuda_available import cuda_available
  from gpu_info import gpu_info
  from vram_budget import vram_budget

  cuda = cuda_available()
  gpus = gpu_info()
  print("CUDA:", cuda)
  print("GPUs:", gpus)

source_raw:
  import sys, os
  FN_ROOT = os.environ.get("FN_REGISTRY_ROOT", "/home/lucas/fn_registry")
  sys.path.insert(0, os.path.join(FN_ROOT, "python/functions/ml"))

  from cuda_available import cuda_available
  from gpu_info import gpu_info
  from vram_budget import vram_budget

  cuda = cuda_available()
  gpus = gpu_info()
  print("CUDA:", cuda)
  print("GPUs:", gpus)


source:
  total_mb = gpus[0]["vram_total_mb"] if gpus else 8000  # fallback CPU/sim
  budget = vram_budget(
      gpu_vram_total_mb=total_mb,
      model_type="sd15",
      quantization="fp16",
      width=512, height=512,
  )
  print("VRAM total:", total_mb, "MB")
  print("Budget:", budget)

source_raw:
  total_mb = gpus[0]["vram_total_mb"] if gpus else 8000  # fallback CPU/sim
  budget = vram_budget(
      gpu_vram_total_mb=total_mb,
      model_type="sd15",
      quantization="fp16",
      width=512, height=512,
  )
  print("VRAM total:", total_mb, "MB")
  print("Budget:", budget)


source:
  # Forzar imports limpios para evitar double-import pydantic
  for _mod in ["sampler_name", "model_ref", "lora_ref", "generation_config", "image_gen_result"]:
      sys.modules.pop(_mod, None)

  from model_ref import ModelRef
  from generation_config import GenerationConfig

  VAULT = "/home/lucas/vaults/imagegen_models/diffusers/sd-turbo"

  model = ModelRef(
      name="stabilityai/sd-turbo",
      model_type="sd15",
      quantization="fp16",
      path=VAULT,
  )

  PROMPTS = [
      ("a cinematic shot of a baby racoon wearing a tiny crown", 42),
      ("watercolor of a cozy library with floating books", 123),
      ("isometric pixel art of a robot fishing on a pier", 7),
      ("oil painting of a fox playing chess against a cat", 999),
  ]

  configs = [
      GenerationConfig(
          prompt=p,
          seed=s,
          steps=1,
          cfg_scale=0.0,
          sampler="euler_a",
          width=512,
          height=512,
          model=model,
      )
      for p, s in PROMPTS
  ]
  for c in configs:
      print(c.seed, c.prompt[:60])

source_raw:
  # Forzar imports limpios para evitar double-import pydantic
  for _mod in ["sampler_name", "model_ref", "lora_ref", "generation_config", "image_gen_result"]:
      sys.modules.pop(_mod, None)

  from model_ref import ModelRef
  from generation_config import GenerationConfig

  VAULT = "/home/lucas/vaults/imagegen_models/diffusers/sd-turbo"

  model = ModelRef(
      name="stabilityai/sd-turbo",
      model_type="sd15",
      quantization="fp16",
      path=VAULT,
  )

  PROMPTS = [
      ("a cinematic shot of a baby racoon wearing a tiny crown", 42),
      ("watercolor of a cozy library with floating books", 123),
      ("isometric pixel art of a robot fishing on a pier", 7),
      ("oil painting of a fox playing chess against a cat", 999),
  ]

  configs = [
      GenerationConfig(
          prompt=p,
          seed=s,
          steps=1,
          cfg_scale=0.0,
          sampler="euler_a",
          width=512,
          height=512,
          model=model,
      )
      for p, s in PROMPTS
  ]
  for c in configs:
      print(c.seed, c.prompt[:60])


source:
  from diffusers_load_pipeline import diffusers_load_pipeline

  pipe = diffusers_load_pipeline(model=model, device="auto", dtype="fp16")
  print("Pipeline:", type(pipe).__name__)
  print("Scheduler:", type(pipe.scheduler).__name__)
  print("Device:", next(pipe.unet.parameters()).device)

source_raw:
  from diffusers_load_pipeline import diffusers_load_pipeline

  pipe = diffusers_load_pipeline(model=model, device="auto", dtype="fp16")
  print("Pipeline:", type(pipe).__name__)
  print("Scheduler:", type(pipe.scheduler).__name__)
  print("Device:", next(pipe.unet.parameters()).device)


source:
  # Forzar imports limpios para alinear tipos con los que usa diffusers_generate
  for _mod in ["image_gen_result", "generation_config", "model_ref", "lora_ref", "sampler_name"]:
      sys.modules.pop(_mod, None)

  from diffusers_generate import diffusers_generate
  from generation_config import GenerationConfig as GC2
  from model_ref import ModelRef as MR2

  # Reconstruir configs con los tipos recien re-importados (alineados con diffusers_generate)
  model2 = MR2(**model.model_dump())
  configs2 = [GC2(**c.model_dump()) for c in configs]

  results = []
  for c in configs2:
      r = diffusers_generate(pipe, c)
      print(f"seed={c.seed} duration={r.duration_ms}ms vram_peak={r.vram_peak_mb}MB")
      results.append(r)
  print("Total:", len(results))

source_raw:
  # Forzar imports limpios para alinear tipos con los que usa diffusers_generate
  for _mod in ["image_gen_result", "generation_config", "model_ref", "lora_ref", "sampler_name"]:
      sys.modules.pop(_mod, None)

  from diffusers_generate import diffusers_generate
  from generation_config import GenerationConfig as GC2
  from model_ref import ModelRef as MR2

  # Reconstruir configs con los tipos recien re-importados (alineados con diffusers_generate)
  model2 = MR2(**model.model_dump())
  configs2 = [GC2(**c.model_dump()) for c in configs]

  results = []
  for c in configs2:
      r = diffusers_generate(pipe, c)
      print(f"seed={c.seed} duration={r.duration_ms}ms vram_peak={r.vram_peak_mb}MB")
      results.append(r)
  print("Total:", len(results))


source:
  from image_grid import image_grid

  labels = [f"seed={r.meta.get('seed', '?')} | {c.prompt[:40]}..." for c, r in zip(configs2, results)]
  grid = image_grid(
      images=[r.image for r in results],
      cols=2,
      labels=labels,
      gap_px=12,
  )
  grid

source_raw:
  from image_grid import image_grid

  labels = [f"seed={r.meta.get('seed', '?')} | {c.prompt[:40]}..." for c, r in zip(configs2, results)]
  grid = image_grid(
      images=[r.image for r in results],
      cols=2,
      labels=labels,
      gap_px=12,
  )
  grid


source:
  from image_save_png import image_save_png
  from genconfig_save_json import genconfig_save_json
  import datetime, os as _os

  OUT_DIR = "/home/lucas/vaults/imagegen_models/outputs"
  CFG_DIR = "/home/lucas/vaults/imagegen_models/configs"
  stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

  grid_path = image_save_png(
      grid,
      _os.path.join(OUT_DIR, f"spike01_grid_{stamp}.png"),
      metadata={"experiment": "spike01_sd_turbo", "model": model2.name, "n": str(len(results))},
  )
  print("grid:", grid_path)

  for i, (c, r) in enumerate(zip(configs2, results)):
      img_path = image_save_png(
          r.image,
          _os.path.join(OUT_DIR, f"spike01_seed{c.seed}_{stamp}.png"),
          metadata={"prompt": c.prompt, "seed": str(c.seed), "steps": str(c.steps), "sampler": c.sampler},
      )
      cfg_path = genconfig_save_json(c, _os.path.join(CFG_DIR, f"spike01_seed{c.seed}_{stamp}.json"))
      print(f"  [{i}] {img_path}")
      print(f"      {cfg_path}")

source_raw:
  from image_save_png import image_save_png
  from genconfig_save_json import genconfig_save_json
  import datetime, os as _os

  OUT_DIR = "/home/lucas/vaults/imagegen_models/outputs"
  CFG_DIR = "/home/lucas/vaults/imagegen_models/configs"
  stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

  grid_path = image_save_png(
      grid,
      _os.path.join(OUT_DIR, f"spike01_grid_{stamp}.png"),
      metadata={"experiment": "spike01_sd_turbo", "model": model2.name, "n": str(len(results))},
  )
  print("grid:", grid_path)

  for i, (c, r) in enumerate(zip(configs2, results)):
      img_path = image_save_png(
          r.image,
          _os.path.join(OUT_DIR, f"spike01_seed{c.seed}_{stamp}.png"),
          metadata={"prompt": c.prompt, "seed": str(c.seed), "steps": str(c.steps), "sampler": c.sampler},
      )
      cfg_path = genconfig_save_json(c, _os.path.join(CFG_DIR, f"spike01_seed{c.seed}_{stamp}.json"))
      print(f"  [{i}] {img_path}")
      print(f"      {cfg_path}")


source:
  rows = []
  for c, r in zip(configs2, results):
      rows.append({
          "seed": c.seed,
          "prompt": c.prompt[:50] + "...",
          "duration_ms": r.duration_ms,
          "vram_peak_mb": r.vram_peak_mb,
      })

  try:
      import pandas as pd
      df = pd.DataFrame(rows)
      display(df)
  except ImportError:
      for row in rows:
          print(row)

  max_dur = max(r.duration_ms for r in results)
  max_vram = max((r.vram_peak_mb or 0) for r in results)
  pass_dur = max_dur < 5000
  pass_vram = max_vram < 6000
  verdict = "PASS" if (pass_dur and pass_vram) else "FAIL"
  print(f"\nMax duration: {max_dur} ms  (target <5000) -> {'OK' if pass_dur else 'FAIL'}")
  print(f"Max VRAM peak: {max_vram} MB  (target <6000) -> {'OK' if pass_vram else 'FAIL'}")
  print(f"\nVEREDICTO: {verdict}")

source_raw:
  rows = []
  for c, r in zip(configs2, results):
      rows.append({
          "seed": c.seed,
          "prompt": c.prompt[:50] + "...",
          "duration_ms": r.duration_ms,
          "vram_peak_mb": r.vram_peak_mb,
      })

  try:
      import pandas as pd
      df = pd.DataFrame(rows)
      display(df)
  except ImportError:
      for row in rows:
          print(row)

  max_dur = max(r.duration_ms for r in results)
  max_vram = max((r.vram_peak_mb or 0) for r in results)
  pass_dur = max_dur < 5000
  pass_vram = max_vram < 6000
  verdict = "PASS" if (pass_dur and pass_vram) else "FAIL"
  print(f"\nMax duration: {max_dur} ms  (target <5000) -> {'OK' if pass_dur else 'FAIL'}")
  print(f"Max VRAM peak: {max_vram} MB  (target <6000) -> {'OK' if pass_vram else 'FAIL'}")
  print(f"\nVEREDICTO: {verdict}")


source:
  from diffusers_unload import diffusers_unload
  diffusers_unload(pipe)
  diffusers_unload(None)  # clear cache global
  print("unloaded")

source_raw:
  from diffusers_unload import diffusers_unload
  diffusers_unload(pipe)
  diffusers_unload(None)  # clear cache global
  print("unloaded")


