File size: 2,730 Bytes
d2a1a41 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Explicit local Quasar inference. No model/codec download or tool execution."""
from pathlib import Path
import argparse
import os
import sys
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint', type=Path, default=ROOT)
parser.add_argument('--task', choices=('text', 'image', 'video'), required=True)
parser.add_argument('--prompt', required=True)
parser.add_argument('--device', choices=('cpu', 'cuda'), default='cuda')
parser.add_argument('--output', type=Path)
parser.add_argument('--resolution', type=int, choices=(64, 128), default=128)
parser.add_argument('--frames', type=int, default=8)
parser.add_argument('--seed', type=int, default=20261021)
parser.add_argument('--max-new-tokens', type=int, default=64)
args = parser.parse_args()
if args.device == 'cpu':
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ.setdefault('TORCHINDUCTOR_TRITON_DISABLE_DEVICE_DETECTION', '1')
# These FP32/BF16 weights do not use the optional GPU quantizer.
sys.modules['torchao'] = None
import torch
torch.set_num_threads(2)
if args.device == 'cuda':
from quasar.training.hardware import check_gpu
check_gpu({'require_cuda': True, 'exclusive_gpu': True, 'min_free_vram_gb': 4})
from quasar.inference.extended import ExtendedGenerator
from quasar.inference.decoding import Sampling
generator = ExtendedGenerator(args.checkpoint.resolve(), args.device)
sampling = Sampling(temperature=.5, top_k=64, top_p=.9, seed=args.seed, max_new_tokens=args.max_new_tokens)
if args.task == 'text':
result = generator.generate_text(args.prompt, sampling)
if args.output:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(result, encoding='utf-8')
print(result)
return
from PIL import Image
output = args.output or Path('generated.png' if args.task == 'image' else 'generated.gif')
output.parent.mkdir(parents=True, exist_ok=True)
if args.task == 'image':
pixels = generator.generate_image(args.prompt, sampling, resolution=args.resolution)
Image.fromarray(pixels).save(output)
else:
if output.suffix.lower() != '.gif':
parser.error('Video preview output must use .gif')
frames = generator.generate_video(args.prompt, args.frames, sampling)
images = [Image.fromarray(frame) for frame in frames]
images[0].save(output, save_all=True, append_images=images[1:], duration=250, loop=0, disposal=2)
print(output.resolve())
if __name__ == '__main__':
main()
|