generate.py

# generate.py
import torch
from generator import Generator128
from torchvision.utils import save_image
import os
import argparse

def main(seed=42, output="laika_dream.png", model_path="model.pth"):
    device = "cpu"
    
    # Загрузка модели
    print("Загрузка LaikaDream...")
    netG = Generator128()
    netG.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
    netG.eval()

    # Генерация
    print(f"Генерация изображения (seed={seed})...")
    torch.manual_seed(seed)
    noise = torch.randn(1, 100, 1, 1, device=device)

    with torch.no_grad():
        fake = netG(noise).detach().cpu()

    # Сохранение: [-1,1] → [0,1]
    save_image(fake * 0.5 + 0.5, output)
    print(f"✅ Изображение сохранено: {os.path.abspath(output)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="LaikaDream — генерация сна Лайки")
    parser.add_argument("--seed", type=int, default=42, help="Случайное семя")
    parser.add_argument("--output", type=str, default="laika_dream.png", help="Путь к выходному файлу")
    parser.add_argument("--model", type=str, default="model.pth", help="Путь к model.pth")
    args = parser.parse_args()
    
    main(seed=args.seed, output=args.output, model_path=args.model)