| """The model behind the demo. |
| |
| Contract used by app.py: |
| MODES {label: mode} shown in the input-type selector |
| GAUSSIANS {label: count} shown in the Gaussian-count dropdown |
| DEFAULTS default settings |
| generate(...) -> (path of the .ply, the image the model saw), called on the GPU worker |
| """ |
| import os |
| import tempfile |
| import uuid |
|
|
| from idol import IdolPipeline |
|
|
| MODES = {'Synthetic': 'synthetic', 'Real photo': 'real'} |
| GAUSSIANS = {'65,536': 65536, '131,072': 131072, '262,144': 262144, '524,288': 524288, '1,048,576': 1048576} |
| DEFAULTS = {'steps': 50, 'cfg': 3.0, 'rescale_t': 3.0, 'gaussians': '262,144', 'seed': 42} |
|
|
| |
| pipeline = IdolPipeline.from_pretrained(os.environ.get('IDOL_WEIGHTS', 'speridlabs/idol'), device='cuda') |
| OUTPUT_DIR = tempfile.mkdtemp(prefix='idol-') |
|
|
|
|
| def generate(image, mode, steps, cfg_strength, rescale_t, seed, num_gaussians): |
| output = pipeline(image, mode=mode, seed=seed, steps=steps, cfg_strength=cfg_strength, rescale_t=rescale_t, |
| num_gaussians=num_gaussians) |
| path = os.path.join(OUTPUT_DIR, f'idol-seed{seed}-{num_gaussians}-{uuid.uuid4().hex[:8]}.ply') |
| output.gaussian.save_ply(path) |
| return path, output.preprocessed_image |
|
|
