Spaces:
Running
on
T4
Running
on
T4
| import os | |
| import shutil | |
| import subprocess | |
| import uuid | |
| import gradio as gr | |
| from pathlib import Path | |
| from typing import Tuple | |
| from app.inference import inference_batch_postprocess | |
| # Should be refactored in the future | |
| class UncondGeneratingMethod(): | |
| def __init__(self, output_main_dir=Path('./outputs')): | |
| self.output_main_dir = output_main_dir | |
| def generate(self): | |
| def generate_uncond(seed, state: gr.BrowserState): | |
| try: | |
| state = check_user_output_dir(state, self.output_main_dir) | |
| generate_output = Path(state['user_output_dir']) / 'unconditional' | |
| os.makedirs(generate_output, exist_ok=True) | |
| if len(os.listdir(generate_output)) > 0: | |
| shutil.rmtree(generate_output) | |
| os.makedirs(generate_output, exist_ok=True) | |
| # Get the generated model | |
| command = [ | |
| "python", "-m", "diffusion.train_diffusion", | |
| "trainer.evaluate=true", | |
| "trainer.batch_size=1000", | |
| "trainer.gpu=1", | |
| f"trainer.test_output_dir={generate_output.as_posix()}", | |
| "trainer.resume_from_checkpoint=YuXingyao/HoLa-Brep/Diffusion_uncond_1100k.ckpt", | |
| "trainer.num_worker=1", | |
| "trainer.accelerator=\"32-true\"", | |
| "trainer.exp_name=test", | |
| "dataset.name=Dummy_dataset", | |
| "dataset.length=32", | |
| "dataset.num_max_faces=30", | |
| "dataset.condition=None", | |
| f"dataset.random_seed={seed}", | |
| "model.name=Diffusion_condition", | |
| "model.autoencoder_weights=YuXingyao/HoLa-Brep/AE_deepcad_1100k.ckpt", | |
| "model.autoencoder=AutoEncoder_1119_light", | |
| "model.with_intersection=true", | |
| "model.in_channels=6", | |
| "model.dim_shape=768", | |
| "model.dim_latent=8", | |
| "model.gaussian_weights=1e-6", | |
| "model.pad_method=random", | |
| "model.diffusion_latent=768", | |
| "model.diffusion_type=epsilon", | |
| "model.gaussian_weights=1e-6", | |
| "model.condition=None", | |
| "model.num_max_faces=30", | |
| "model.beta_schedule=linear", | |
| "model.addition_tag=false", | |
| "model.name=Diffusion_condition" | |
| ] | |
| env = os.environ.copy() | |
| env["CUDA_VISIBLE_DEVICES"] = "0" | |
| gr.Info("Start diffusing", title="Runtime Info") | |
| subprocess.run(command, check=True, env=env) | |
| gr.Info("Finished diffusing", title="Runtime Info") | |
| # Post-process the generated model | |
| postprocess_output = Path(state['user_output_dir']) / 'unconditional_post' | |
| os.makedirs(postprocess_output, exist_ok=True) | |
| if len(os.listdir(postprocess_output)) > 0: | |
| shutil.rmtree(postprocess_output) | |
| os.makedirs(postprocess_output, exist_ok=True) | |
| gr.Info("Start post-processing.", title="Runtime Info") | |
| inference_batch_postprocess( | |
| file_dir=generate_output.as_posix(), | |
| output_dir=postprocess_output.as_posix(), | |
| num_cpus=2, | |
| drop_num=0 | |
| ) | |
| gr.Info("Finished post-processing!", title="Runtime Info") | |
| valid_models = get_valid_models(postprocess_output) | |
| # Should have valid outputs | |
| if len(valid_models) <= 0: | |
| raise UncondGeneraingException("No Valid Model Generated!") | |
| # Update the user state | |
| state["uncond"] = list() | |
| for i, model_number in enumerate(valid_models): | |
| if (postprocess_output / model_number / 'debug_face_loop' / 'optimized_edge.obj').exists(): | |
| edge = (postprocess_output / model_number / 'debug_face_loop' / 'optimized_edge.obj').as_posix() | |
| else: | |
| edge = (postprocess_output / model_number / 'debug_face_loop' / 'edge.obj').as_posix() # Hard coding is not good. | |
| solid = (postprocess_output / model_number / 'recon_brep.stl').as_posix() | |
| step = (postprocess_output / model_number / 'recon_brep.step').as_posix() | |
| state["uncond"].append([edge, solid, step]) | |
| gr.Info(f"{len(valid_models) if len(valid_models) < 4 else 4} valid models generated!", title="Finish generating") | |
| edge_file = state["uncond"][0][0] | |
| solid_file = state["uncond"][0][1] | |
| step_file = state["uncond"][0][2] | |
| return edge_file, solid_file, step_file, state["uncond"][0], state | |
| except UncondEmptyInputException as input_e: | |
| gr.Warning(str(input_e), title="Empty Input") | |
| except UncondGeneraingException as generating_e: | |
| gr.Warning(str(generating_e), title="No Valid Generation") | |
| except Exception as e: | |
| print(e) | |
| gr.Warning("Something bad happened. Please try some other models", title="Unknown Error") | |
| return gr.update(), gr.update(), gr.update(), gr.update(), state | |
| return generate_uncond | |
| def get_valid_models(postprocess_output: Path) -> Tuple[Path, Path, Path]: | |
| output_folders = [model_folder for model_folder in os.listdir(postprocess_output) if 'success.txt' in os.listdir(postprocess_output / model_folder)] | |
| return output_folders | |
| def check_user_output_dir(state: dict, output_dir): | |
| if state['user_id'] is None: | |
| state['user_id'] = uuid.uuid4() | |
| if state['user_output_dir'] is None: | |
| state['user_output_dir'] = Path(output_dir) / f"user_{state['user_id']}" | |
| os.makedirs(state['user_output_dir'], exist_ok=True) | |
| return state | |
| class UncondGeneraingException(Exception): | |
| """Custom exception if generating failed.""" | |
| pass | |
| class UncondEmptyInputException(Exception): | |
| """Custom exception if the input is empty.""" | |
| pass |