Spaces:
Sleeping
Sleeping
| """Strict SkinTokens worker: prewarmed model, separate Blender process, preserved input.""" | |
| from pathlib import Path | |
| import hashlib | |
| import json | |
| import shutil | |
| import uuid | |
| import time | |
| import traceback | |
| import spaces | |
| import gradio as gr | |
| import torch | |
| from prepare import prepare, LOCK, ROOT | |
| from preservation import verify | |
| prepare() | |
| from src.data.dataset import DatasetConfig, RigDatasetModule # noqa: E402 | |
| from src.data.transform import Transform # noqa: E402 | |
| from src.tokenizer.parse import get_tokenizer # noqa: E402 | |
| from src.server.spec import BpyClient, bpy_server_request, get_model # noqa: E402 | |
| # Blender CPU imports and all CUDA model placement happen before a user books a GPU. | |
| BPY = BpyClient.launch() | |
| MODEL = get_model(LOCK['checkpoint']).eval() | |
| TOKENIZER = get_tokenizer(**MODEL.tokenizer_config) | |
| TRANSFORM = Transform.parse(**MODEL.transform_config['predict_transform']) | |
| def sha256(path): | |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() | |
| def predict(source): | |
| config = DatasetConfig.parse( | |
| shuffle=False, batch_size=1, num_workers=0, pin_memory=False, persistent_workers=False, | |
| datapath={'data_name': None, 'loader': 'bpy_server', 'filepaths': {'articulation': [str(source)]}}, | |
| ).split_by_cls() | |
| module = RigDatasetModule(predict_dataset_config=config, predict_transform=TRANSFORM, | |
| tokenizer=TOKENIZER, process_fn=MODEL._process_fn) | |
| batch = next(iter(module.predict_dataloader()['articulation'])) | |
| batch = {key: value.to(MODEL.device) if isinstance(value, torch.Tensor) else value | |
| for key, value in batch.items()} | |
| batch.pop('skeleton_tokens', None) | |
| batch.pop('skeleton_mask', None) | |
| batch['generate_kwargs'] = dict(max_length=2048, top_k=5, top_p=.95, temperature=1.0, | |
| repetition_penalty=2.0, num_return_sequences=1, num_beams=10, do_sample=True) | |
| with torch.inference_mode(): | |
| asset = MODEL.predict_step(batch, skeleton_tokens=None, make_asset=True)['results'][0].asset | |
| if asset is None or asset.skin is None: | |
| raise RuntimeError('SkinTokens produced no skinned asset') | |
| # Asset stores NumPy data; no CUDA tensors leave this worker. | |
| return asset | |
| def rig(source): | |
| global BPY | |
| output = ROOT / 'outputs' / ('rig-' + uuid.uuid4().hex) | |
| output.mkdir(parents=True) | |
| original = output / 'source.glb' | |
| shutil.copyfile(source, original) | |
| receipt = {'provider': 'native SkinTokens', 'lock': LOCK, 'source_sha256': sha256(original), | |
| 'started': time.time(), 'stage': 'load', 'status': 'running'} | |
| try: | |
| # Surface a Blender/input failure before requesting the GPU. | |
| if not BPY.is_alive(): | |
| BPY = BpyClient.launch() | |
| bpy_server_request('load', str(original)) | |
| receipt['stage'] = 'predict' | |
| asset = predict(original) | |
| receipt['stage'] = 'transfer' | |
| result = output / 'rigged.glb' | |
| transfer = bpy_server_request('transfer', dict(source_asset=asset, target_path=str(original), | |
| export_path=str(result), group_per_vertex=4)) | |
| if transfer != 'ok': | |
| raise RuntimeError(f'Blender transfer failed: {transfer}') | |
| receipt['stage'] = 'verify' | |
| receipt['preservation'] = verify(original, result) | |
| receipt.update(status='complete', output_sha256=sha256(result), finished=time.time()) | |
| return str(result), receipt | |
| except Exception as error: | |
| traceback.print_exc() | |
| message = f"{receipt['stage']}: {type(error).__name__}: {error}" | |
| receipt.update(status='failed', error=message, finished=time.time()) | |
| raise gr.Error(message) from error | |
| finally: | |
| (output / 'provenance.json').write_text(json.dumps(receipt, indent=2)) | |
| with gr.Blocks(title='Component Studio SkinTokens') as demo: | |
| source = gr.File(file_types=['.glb'], type='filepath', label='Textured GLB') | |
| button = gr.Button('Rig character') | |
| result = gr.File(label='Rigged GLB') | |
| report = gr.JSON(label='Preservation and provenance') | |
| button.click(rig, source, [result, report], api_name='rig', concurrency_limit=1) | |
| demo.queue(max_size=4).launch(server_name='0.0.0.0', server_port=7860, | |
| allowed_paths=[str(ROOT / 'outputs')]) | |