| """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') |
| |
| 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() |
|
|
