# Project export: Aira This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: TreeHacks 2024 - Tagline: iOS app that generates personalised AR tutorials using 3D stable diffusion in real time. - Devpost: https://devpost.com/software/aira-9sqo1j - GitHub: https://github.com/nkoorty/TreeHacks - Video: https://www.youtube.com/embed/51pbTIhZ5MA?enablejsapi=1&hl=en_US&rel=0&start=&version=3&wmode=transparent - Team: 3 GitHub contributor(s) — jjjutla (13 commits), Artemiy (3 commits), Adesh Dooraree (1 commits) ## Devpost submission (written by the team) ### Inspiration Not all products are designed in a user-friendly and intuitive way. We often come across devices that are annoying and unclear to use. This is especially true for people with less exposure to tech, such as seniors. Whether it’s setting up a new tech gadget or controlling the AC in a new rental car, reading long user manuals or finding a random YouTube tutorial is currently the best course of action. But what if an AI could generate the tutorial specifically for you directly on your phone and visually explain the product using interactive AR? ### What it does We leave AI chatbots in the dust, by combining them with 3D stable diffusion and Augmented Reality to create a user experience as if an expert is physically next to you, visually answering your question with a helpful virtual demonstration. Workflow User wants to know how to interact with an object. They open the app and place their camera in-front of the object. The user asks their question e.g. How do I do 'X'? Object detection model detects the item in-front of user. Speech to text understands the user’s question and sends the label and prompt to the backend LLM instruction agent. The instruction agent takes the user's prompt and generates a list of clear instructions to resolve the user’s problem. The detected object and contextualised instructions are fed into a 3D stable diffusion model which generates a digital twin that is displayed alongside the real object in AR. The 3D models are positioned in AR space as a visual guidance for the written instructions, which are also shown to the user. ### How we built it FrontEnd: The core frontend was developed using Swift UI, using ARKit for rendering the tutorials in space and CoreML as the on-device model to detect the object in front of the camera. We also used AVFoundation to enable speech-to-text capabilities to simplify the user experience. For more complex and involved tutorials we aim to make the frontend compatible with the Apple Vision Pro in the near future. Instruction Agent: The instruction agent simplifies user guidance by generating concise instructions in three clear steps. It receives prompts via a REST API from the front-end, incorporating them into the output JSON format. These instructions are then contextualised for the Text-to-3D model, which facilitates the generation and positioning of AR objects. This process involves passing the question and label through a LLM to produce the finalised JSON. Text to 3D Stable Diffusion: The text to 3D stable diffusion model was developed using a pre-trained 2D text-to-image diffusion model to perform text-to-3D synthesis. We used probability density distillation loss to optimise a NeRF model using gradient descent. The resulting model can be viewed from any angle and requires no 3D training data or modifications to the image diffusion model. Because querying each ray in a NeRF requires a lot of computation we used a Sparse Neural Radiance Grid (SNeRG) that enables real-time rendering. This involved reformulation of the architecture using a sparse voxel grid representation with learned feature vectors. We used USDPython with ARConvert for Usdz compatibility on iOS. The following papers were used as technical support and inspiration: https://instruct-nerf2nerf.github.io/ https://phog.github.io/snerg/ https://dreamfusion3d.github.io/ ### Challenges we ran into Rendering 3D models at high speed and quality turned out to be very tough. Our model started out producing low quality AR objects within one minute, and after precomputing and storing the NeRF into a SNeRG, we were able to cut that time down to several seconds. Producing the highest quality models takes longer and is a challenge that we want to address in the future. For now, the lower quality version suffices, and on the small size of a smartphone screen, is not much of an issue. ### Accomplishments we're proud of We made a fully functional demo and MVP! Despite facing many technical challenges along the way, we managed to overcome them all and are proud of the functionality and complexity of our product. We were able to integrate many packages and models into a complex pipeline that seamlessly converts the user’s question into a visual tutorial. The technical complexity of our solution was both challenging and rewarding, and we are excited to work on this further and see how far we can push the performance and quality of the model, especially considering how close to the edge of research it is. ### What we learned We used many new packages and techniques in this project, significantly expanding our skillset. Our biggest breakthrough was getting the 3D stable diffusion algorithm to work, as this was something we had never done before. We also expanded our AR capabilities by learning about ARKit, RealityKit and AVFoundation as well as using the ‘Combine’ and ‘Speech’ packages to transcribe the user’s spoken prompt and ensure a smooth experience. ### What's next Our next goal is to improve the model to animate the AR objects generated using 3D stable diffusion. This involves identifying each moving component as a separate object, generating them separately, and then using the contextualisation ability of the instruction agent to understand the desired movement of the components relative to each other, and outputting the motion in polar coordinates. Following this, we will further fine-tune and optimise our model to cut down the time it takes to generate the 3D AR models. To improve the UX we also plan to add arrows visualising the actions needed to interact with the user too. Deck: https://www.canva.com/design/DAF9EZRlAW8/lDw9k8mMUDGqLUeVQBfBbw/edit?utm_content=DAF9EZRlAW8&utm_campaign=designshare&utm_medium=link2&utm_source=sharebutton ## README (from the GitHub repository) # Aira - AI Generated AR Tutorials **Aira is an iOS app that generates AR tutorials using 3D stable diffusion in real time.** **Devpost: https://devpost.com/software/aira-9sqo1j** ## Inspiration Not all products are designed in a user-friendly and intuitive way. We often come across devices that are annoying and unclear to use. This is especially true for people with less exposure to tech, such as seniors. Whether it’s setting up a new tech gadget or controlling the AC in a new rental car, reading long user manuals or finding a random YouTube tutorial is currently the best course of action. But what if an AI could generate the tutorial specifically for you directly on your phone and visually explain the product using interactive AR? ## Images: thumbnail Demo Turtle camera tutorial ## Workflow 1. User wants to know how to interact with an object. 2. They open the app and place their camera in-front of the object. 3. The user asks their question e.g. How do I do 'X'? 4. Object detection model detects the item in-front of user. 5. Speech to text understands the user’s question and sends the label and prompt to the backend LLM instruction agent. 6. The instruction agent takes the user's prompt and generates a list of clear instructions to resolve the user’s problem. 7. The detected object and contextualised instructions are fed into a 3D stable diffusion model which generates a digital twin that is displayed alongside the real object in AR. 8. The 3D models are positioned in AR space as a visual guidance for the written instructions, which are also shown to the user. ## How we built it **FrontEnd:** The core frontend was developed using Swift UI, using ARKit for rendering the tutorials in space and CoreML as the on-device model to detect the object in front of the camera. We also used AVFoundation to enable speech-to-text capabilities to simplify the user experience. For more complex and involved tutorials we aim to make the frontend compatible with the Apple Vision Pro in the near future. **Instruction Agent:** The instruction agent simplifies user guidance by generating concise instructions in three clear steps. It receives prompts via a REST API from the front-end, incorporating them into the output JSON format. These instructions are then contextualised for the Text-to-3D model, which facilitates the generation and positioning of AR objects. This process involves passing the question and label through a LLM to produce the finalised JSON. **Text to 3D Stable Diffusion:** The text to 3D stable diffusion model was developed using a pre-trained 2D text-to-image diffusion model to perform text-to-3D synthesis. We used probability density distillation loss to optimise a NeRF model using gradient descent. The resulting model can be viewed from any angle and requires no 3D training data or modifications to the image diffusion model. Because querying each ray in a NeRF requires a lot of computation we used a Sparse Neural Radiance Grid (SNeRG) that enables real-time rendering. This involved reformulation of the architecture using a sparse voxel grid representation with learned feature vectors. We used USDPython with ARConvert for Usdz compatibility on iOS. The following papers were used as technical support and inspiration: - https://instruct-nerf2nerf.github.io/ - https://phog.github.io/snerg/ - https://dreamfusion3d.github.io/ ## License MIT ## Detected evidence (automated analysis) Indexed codebase: 51 recognized source files, 357 KB. - C (language) — detected in the code - C++ (language) — detected in the code - Hugging Face (technology) — detected in the code - Python (language) — detected in the code - PyTorch (technology) — detected in the code - Swift (language) — detected in the code - Rust (language) — claimed on Devpost, not found in the code - TypeScript (language) — claimed on Devpost, not found in the code ## Codebase structure (from repository index) ### Files (84 of 84) ``` .DS_Store LICENSE README.md server/.DS_Store server/.gitignore server/app.py server/latenighttest.py stablediffusion/.DS_Store stablediffusion/.gitignore stablediffusion/activation.py stablediffusion/dpt.py stablediffusion/encoding.py stablediffusion/evaluation/mesh_to_video.py stablediffusion/evaluation/Prompt.py stablediffusion/evaluation/r_precision.py stablediffusion/evaluation/readme.md stablediffusion/freqencoder/__init__.py stablediffusion/freqencoder/backend.py stablediffusion/freqencoder/freq.py stablediffusion/freqencoder/setup.py stablediffusion/freqencoder/src/bindings.cpp stablediffusion/freqencoder/src/freqencoder.cu stablediffusion/freqencoder/src/freqencoder.h stablediffusion/gridencoder/__init__.py stablediffusion/gridencoder/backend.py stablediffusion/gridencoder/grid.py stablediffusion/gridencoder/setup.py stablediffusion/gridencoder/src/bindings.cpp stablediffusion/gridencoder/src/gridencoder.cu stablediffusion/gridencoder/src/gridencoder.h stablediffusion/guidance/clip_utils.py stablediffusion/guidance/if_utils.py stablediffusion/guidance/perpneg_utils.py stablediffusion/guidance/sd_utils.py stablediffusion/guidance/zero123_utils.py stablediffusion/main.py stablediffusion/meshutils.py stablediffusion/nerf/gui.py stablediffusion/nerf/network_grid_taichi.py stablediffusion/nerf/network_grid_tcnn.py stablediffusion/nerf/network_grid.py stablediffusion/nerf/network.py stablediffusion/nerf/provider.py stablediffusion/nerf/renderer.py stablediffusion/nerf/utils.py stablediffusion/optimizer.py stablediffusion/preprocess_image.py stablediffusion/raymarching/__init__.py stablediffusion/raymarching/.DS_Store stablediffusion/raymarching/backend.py stablediffusion/raymarching/raymarching.py stablediffusion/raymarching/setup.py stablediffusion/raymarching/src/bindings.cpp stablediffusion/raymarching/src/raymarching.cu stablediffusion/raymarching/src/raymarching.h stablediffusion/requirements.txt TreeHacksMobile/.DS_Store TreeHacksMobile/TreeHacks.xcodeproj/project.pbxproj TreeHacksMobile/TreeHacks.xcodeproj/project.xcworkspace/contents.xcworkspacedata TreeHacksMobile/TreeHacks.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist TreeHacksMobile/TreeHacks.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved TreeHacksMobile/TreeHacks.xcodeproj/project.xcworkspace/xcuserdata/artemiym.xcuserdatad/UserInterfaceState.xcuserstate TreeHacksMobile/TreeHacks.xcodeproj/xcuserdata/artemiym.xcuserdatad/xcdebugger/Breakpoints_v2.xcbkptlist TreeHacksMobile/TreeHacks.xcodeproj/xcuserdata/artemiym.xcuserdatad/xcschemes/xcschememanagement.plist TreeHacksMobile/TreeHacks/AppDelegate.swift TreeHacksMobile/TreeHacks/ARModels/ai2.usdz TreeHacksMobile/TreeHacks/ARModels/airpods2.usdz TreeHacksMobile/TreeHacks/ARModels/airpodsbad.usdz TreeHacksMobile/TreeHacks/ARModels/Cosmonaut.reality TreeHacksMobile/TreeHacks/ARModels/man.usdz TreeHacksMobile/TreeHacks/ARModels/tree.usdz TreeHacksMobile/TreeHacks/Assets.xcassets/AccentColor.colorset/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/AppIcon.appiconset/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/AR Resources.arresourcegroup/airpods.arreferenceimage/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/AR Resources.arresourcegroup/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/background.imageset/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/Contents.json TreeHacksMobile/TreeHacks/Assets.xcassets/man.imageset/Contents.json TreeHacksMobile/TreeHacks/ContentView.swift TreeHacksMobile/TreeHacks/ExpandableEvent.swift TreeHacksMobile/TreeHacks/HomeView.swift TreeHacksMobile/TreeHacks/ImageHandler.swift TreeHacksMobile/TreeHacks/Preview Content/Preview Assets.xcassets/Contents.json TreeHacksMobile/TreeHacks/SpeechRecogniser.swift ``` ### Dependencies - stablediffusion/requirements.txt: accelerate, carvekit-colab, dearpygui, debugpy-run, diffusers@>= 0.9.0, einops, gdown, git@+https://github.com/openai/CLIP.git, huggingface_hub, imageio, imageio-ffmpeg, kornia, matplotlib, ninja, numpy, omegaconf, opencv-python, pandas, PyMCubes, pymeshlab, pytorch-lightning, rich, scikit-learn, scipy, sentencepiece, taming-transformers-rom1504, tensorboard, tensorboardX, timm, torch, torch-ema, tqdm, transformers, trimesh, xatlas ### Recent commits (newest first) - Update README.md - ContentView bug fix - Update README.md - Update README.md - Update README.md - Update README.md - Update README.md - Update README.md - Update README.md - change - added iOS code - Update README.md - Update README.md - Update README.md - Update README.md - core model - added server - Update README.md - Update README.md - Initial commit ## Key source files (fetched from GitHub, selected and truncated for size) ### stablediffusion/requirements.txt ``` tqdm rich ninja numpy pandas scipy scikit-learn matplotlib opencv-python imageio imageio-ffmpeg torch torch-ema einops tensorboard tensorboardX # for gui dearpygui # for grid_tcnn # git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch # for stable-diffusion huggingface_hub diffusers >= 0.9.0 accelerate transformers # for dmtet and mesh export xatlas trimesh PyMCubes pymeshlab git+https://github.com/NVlabs/nvdiffrast/ # for zero123 carvekit-colab omegaconf pytorch-lightning taming-transformers-rom1504 kornia git+https://github.com/openai/CLIP.git # for omnidata gdown # for dpt timm # for remote debugging debugpy-run # for deepfloyd if sentencepiece ``` ### server/app.py ```python from flask import Flask, request, jsonify from openai import OpenAI, OpenAIError import subprocess from dotenv import load_dotenv import os load_dotenv() app = Flask(__name__) openai_api_key = os.getenv("OPENAI_API_KEY") openai_client = OpenAI(api_key=openai_api_key) @app.route('/chat', methods=['POST']) def chat(): try: data = request.json question = data.get('prompt', '') label = data.get('label', '') print(type(label)) prompt = f"Write a tutorial on {question} ({label}) aimed at senior citizens who are complete beginners in the subject. Keep the tutorial only 3 lines long - extremely concise to cover only the basics. Ensure that each step is clear and easy to follow. Consider using simple language and providing visual aids where necessary to aid understanding. Your goal is to empower seniors to grasp the fundamental concepts without overwhelming them with technical jargon or advanced techniques. Format your tutorial steps in a JSON (JavaScript Object Notation) structure for clarity and organization." if not prompt: return jsonify({'error': 'Prompt not provided'}), 400 response = openai_client.chat.completions.create( model="gpt-3.5-turbo", messages=[{"role": "user", "content": prompt}] ) gptRaw = response.choices[0].message.content # Path to your shell script script_path = '/Users/adesh/Documents/treehacksGPT/t23d.sh' process = subprocess.Popen(script_path, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) # Input data into the 3D generative AI model input_data = label process.stdin.write(input_data) process.stdin.flush() # buffer output, errors = process.communicate() print("Output:", output) return gptRaw except OpenAIError as e: return jsonify({'error': str(e)}), 500 if __name__ == '__main__': app.run(debug=True, port=3000, host='0.0.0.0') ``` ### stablediffusion/main.py ```python import torch import argparse import pandas as pd import sys from nerf.provider import NeRFDataset from nerf.utils import * if __name__ == '__main__': # See https://stackoverflow.com/questions/27433316/how-to-get-argparse-to-read-arguments-from-a-file-with-an-option-rather-than-pre class LoadFromFile (argparse.Action): def __call__ (self, parser, namespace, values, option_string = None): with values as f: # parse arguments in the file and store them in the target namespace parser.parse_args(f.read().split(), namespace) parser = argparse.ArgumentParser() parser.add_argument('--file', type=open, action=LoadFromFile, help="specify a file filled with more arguments") parser.add_argument('--text', default=None, help="text prompt") parser.add_argument('--negative', default='', type=str, help="negative text prompt") parser.add_argument('-O', action='store_true', help="equals --fp16 --cuda_ray") parser.add_argument('-O2', action='store_true', help="equals --backbone vanilla") parser.add_argument('--test', action='store_true', help="test mode") parser.add_argument('--six_views', action='store_true', help="six_views mode: save the images of the six views") parser.add_argument('--eval_interval', type=int, default=1, help="evaluate on the valid set every interval epochs") parser.add_argument('--test_interval', type=int, default=100, help="test on the test set every interval epochs") parser.add_argument('--workspace', type=str, default='workspace') parser.add_argument('--seed', default=None) parser.add_argument('--image', default=None, help="image prompt") parser.add_argument('--image_config', default=None, help="image config csv") parser.add_argument('--known_view_interval', type=int, default=4, help="train default view with RGB loss every & iters, only valid if --image is not None.") parser.add_argument('--IF', action='store_true', help="experimental: use DeepFloyd IF as the guidance model for nerf stage") parser.add_argument('--guidance', type=str, nargs='*', default=['SD'], help='guidance model') parser.add_argument('--guidance_scale', type=float, default=100, help="diffusion model classifier-free guidance scale") parser.add_argument('--save_mesh', action='store_true', help="export an obj mesh with texture") parser.add_argument('--mcubes_resolution', type=int, default=256, help="mcubes resolution for extracting mesh") parser.add_argument('--decimate_target', type=int, default=5e4, help="target face number for mesh decimation") parser.add_argument('--dmtet', action='store_true', help="use dmtet finetuning") parser.add_argument('--tet_grid_size', type=int, default=128, help="tet grid size") parser.add_argument('--init_with', type=str, default='', help="ckpt to init dmtet") parser.add_argument('--lock_geo', action='store_true', help="disable dmtet to learn geometry") ## Perp-Neg options parser.add_argument('--perpneg', action='store_true', help="use perp_neg") parser.add_argument('--negative_w', type=float, default=-2, help="The scale of the weights of negative prompts. A larger value will help to avoid the Janus problem, but may cause flat faces. Vary between 0 to -4, depending on the prompt") parser.add_argument('--front_decay_factor', type=float, default=2, help="decay factor for the front prompt") parser.add_argument('--side_decay_factor', type=float, default=10, help="decay factor for the side prompt") ### training options parser.add_argument('--iters', type=int, default=10000, help="training iters") parser.add_argument('--lr', type=float, default=1e-3, help="max learning rate") parser.add_argument('--ckpt', type=str, default='latest', help="possible options are ['latest', 'scratch', 'best', 'latest_model']") parser.add_argument('--cuda_ray', action='store_true', help="use CUDA raymarching instead of pytorch") parser.add_argument('--taichi_ray', action='store_true', help="use taichi raymarching") parser.add_argument('--max_steps', type=int, default=1024, help="max num steps sampled per ray (only valid when using --cuda_ray)") parser.add_argument('--num_steps', type=int, default=64, help="num steps sampled per ray (only valid when not using --cuda_ray)") parser.add_argument('--upsample_steps', type=int, default=32, help="num steps up-sampled per ray (only valid when not using --cuda_ray)") parser.add_argument('--update_extra_interval', type=int, default=16, help="iter interval to update extra status (only valid when using --cuda_ray)") parser.add_argument('--max_ray_batch', type=int, default=4096, help="batch size of rays at inference to avoid OOM (only valid when not using --cuda_ray)") parser.add_argument('--latent_iter_ratio', type=float, default=0.2, help="training iters that only use albedo shading") parser.add_argument('--albedo_iter_ratio', type=float, default=0, help="training iters that only use albedo shading") parser.add_argument('--min_ambient_ratio', type=float, default=0.1, help="minimum ambient ratio to use in lambertian shading") parser.add_argument('--textureless_ratio', type=float, default=0.2, help="ratio of textureless shading") parser.add_argument('--jitter_pose', action='store_true', help="add jitters to the randomly sampled camera poses") parser.add_argument('--jitter_center', type=float, default=0.2, help="amount of jitter to add to sampled camera pose's center (camera location)") parser.add_argument('--jitter_target', type=float, default=0.2, help="amount of jitter to add to sampled camera pose's target (i.e. 'look-at')") parser.add_argument('--jitter_up', type=float, default=0.02, help="amount of jitter to add to sampled camera pose's up-axis (i.e. 'camera roll')") parser.add_argument('--uniform_sphere_rate', type=float, default=0, help="likelihood of sampling camera location uniformly on the sphere surfa [truncated — 17258 more characters] ``` ### stablediffusion/activation.py ```python import torch from torch.autograd import Function from torch.cuda.amp import custom_bwd, custom_fwd class _trunc_exp(Function): @staticmethod @custom_fwd(cast_inputs=torch.float) def forward(ctx, x): ctx.save_for_backward(x) return torch.exp(x) @staticmethod @custom_bwd def backward(ctx, g): x = ctx.saved_tensors[0] return g * torch.exp(x.clamp(max=15)) trunc_exp = _trunc_exp.apply def biased_softplus(x, bias=0): return torch.nn.functional.softplus(x - bias) ``` ### server/latenighttest.py ```python from flask import Flask, request, jsonify app = Flask(__name__) @app.route('/submit', methods=['POST']) def submit(): data = request.get_json() new_data = { "1": "Plug in the air fryer and preheat it to the recommended temperature.", "2": "Add your food into the basket, making sure not to overcrowd it.", "3": "Set the time and temperature according to the recipe, then press start to begin cooking.", "model": "tree" } return jsonify(new_data) if __name__ == '__main__': app.run(debug=True, port=3000, host='0.0.0.0') ``` ### stablediffusion/meshutils.py ```python import numpy as np import pymeshlab as pml def poisson_mesh_reconstruction(points, normals=None): import open3d as o3d pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(points) pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10) if normals is None: pcd.estimate_normals() else: pcd.normals = o3d.utility.Vector3dVector(normals[ind]) o3d.visualization.draw_geometries([pcd], point_show_normal=False) mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9) vertices_to_remove = densities < np.quantile(densities, 0.1) mesh.remove_vertices_by_mask(vertices_to_remove) o3d.visualization.draw_geometries([mesh]) vertices = np.asarray(mesh.vertices) triangles = np.asarray(mesh.triangles) print(f'[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}') return vertices, triangles def decimate_mesh(verts, faces, target, backend='pymeshlab', remesh=False, optimalplacement=True): _ori_vert_shape = verts.shape _ori_face_shape = faces.shape if backend == 'pyfqmr': import pyfqmr solver = pyfqmr.Simplify() solver.setMesh(verts, faces) solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False) verts, faces, normals = solver.getMesh() else: m = pml.Mesh(verts, faces) ms = pml.MeshSet() ms.add_mesh(m, 'mesh') # will copy! ms.meshing_decimation_quadric_edge_collapse(targetfacenum=int(target), optimalplacement=optimalplacement) if remesh: ms.meshing_isotropic_explicit_remeshing(iterations=3, targetlen=pml.Percentage(1)) # extract mesh m = ms.current_mesh() verts = m.vertex_matrix() faces = m.face_matrix() print(f'[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}') return verts, faces def clean_mesh(verts, faces, v_pct=1, min_f=8, min_d=5, repair=True, remesh=True, remesh_size=0.01): _ori_vert_shape = verts.shape _ori_face_shape = faces.shape m = pml.Mesh(verts, faces) ms = pml.MeshSet() ms.add_mesh(m, 'mesh') # will copy! # filters ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces if v_pct > 0: ms.meshing_merge_close_vertices(threshold=pml.Percentage(v_pct)) # 1/10000 of bounding box diagonal ms.meshing_remove_duplicate_faces() # faces defined by the same verts ms.meshing_remove_null_faces() # faces with area == 0 if min_d > 0: ms.meshing_remove_connected_component_by_diameter(mincomponentdiag=pml.Percentage(min_d)) if min_f > 0: ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f) if repair: ms.meshing_repair_non_manifold_edges(method=0) ms.meshing_repair_non_manifold_vertices(vertdispratio=0) if remesh: # ms.apply_coord_taubin_smoothing() ms.meshing_isotropic_explicit_remeshing(iterations=3, targetlen=pml.AbsoluteValue(remesh_size)) # extract mesh m = ms.current_mesh() verts = m.vertex_matrix() faces = m.face_matrix() print(f'[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}') return verts, faces ``` ### stablediffusion/encoding.py ```python import torch import torch.nn as nn import torch.nn.functional as F class FreqEncoder_torch(nn.Module): def __init__(self, input_dim, max_freq_log2, N_freqs, log_sampling=True, include_input=True, periodic_fns=(torch.sin, torch.cos)): super().__init__() self.input_dim = input_dim self.include_input = include_input self.periodic_fns = periodic_fns self.N_freqs = N_freqs self.output_dim = 0 if self.include_input: self.output_dim += self.input_dim self.output_dim += self.input_dim * N_freqs * len(self.periodic_fns) if log_sampling: self.freq_bands = 2 ** torch.linspace(0, max_freq_log2, N_freqs) else: self.freq_bands = torch.linspace(2 ** 0, 2 ** max_freq_log2, N_freqs) self.freq_bands = self.freq_bands.numpy().tolist() def forward(self, input, max_level=None, **kwargs): if max_level is None: max_level = self.N_freqs else: max_level = int(max_level * self.N_freqs) out = [] if self.include_input: out.append(input) for i in range(max_level): freq = self.freq_bands[i] for p_fn in self.periodic_fns: out.append(p_fn(input * freq)) # append 0 if self.N_freqs - max_level > 0: out.append(torch.zeros(*input.shape[:-1], (self.N_freqs - max_level) * 2 * input.shape[-1], device=input.device, dtype=input.dtype)) out = torch.cat(out, dim=-1) return out def get_encoder(encoding, input_dim=3, multires=6, degree=4, num_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=2048, align_corners=False, interpolation='linear', **kwargs): if encoding == 'None': return lambda x, **kwargs: x, input_dim elif encoding == 'frequency_torch': encoder = FreqEncoder_torch(input_dim=input_dim, max_freq_log2=multires-1, N_freqs=multires, log_sampling=True) elif encoding == 'frequency': # CUDA implementation, faster than torch. from freqencoder import FreqEncoder encoder = FreqEncoder(input_dim=input_dim, degree=multires) elif encoding == 'sphere_harmonics': from shencoder import SHEncoder encoder = SHEncoder(input_dim=input_dim, degree=degree) elif encoding == 'hashgrid': from gridencoder import GridEncoder encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='hash', align_corners=align_corners, interpolation=interpolation) elif encoding == 'tiledgrid': from gridencoder import GridEncoder encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='tiled', align_corners=align_corners, interpolation=interpolation) elif encoding == 'hashgrid_taichi': from taichi_modules.hash_encoder import HashEncoderTaichi encoder = HashEncoderTaichi(batch_size=4096) #TODO: hard encoded batch size else: raise NotImplementedError('Unknown encoding mode, choose from [None, frequency, sphere_harmonics, hashgrid, tiledgrid]') return encoder, encoder.output_dim ``` ### stablediffusion/preprocess_image.py ```python import os import sys import cv2 import argparse import numpy as np import matplotlib.pyplot as plt import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms from PIL import Image class BackgroundRemoval(): def __init__(self, device='cuda'): from carvekit.api.high import HiInterface self.interface = HiInterface( object_type="object", batch_size_seg=5, batch_size_matting=1, device=device, seg_mask_size=640, matting_mask_size=2048, trimap_prob_threshold=231, trimap_dilation=30, trimap_erosion_iters=5, fp16=True, ) @torch.no_grad() def __call__(self, image): image = Image.fromarray(image) image = self.interface([image])[0] image = np.array(image) return image class BLIP2(): def __init__(self, device='cuda'): self.device = device from transformers import AutoProcessor, Blip2ForConditionalGeneration self.processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16).to(device) @torch.no_grad() def __call__(self, image): image = Image.fromarray(image) inputs = self.processor(image, return_tensors="pt").to(self.device, torch.float16) generated_ids = self.model.generate(**inputs, max_new_tokens=20) generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() return generated_text class DPT(): def __init__(self, task='depth', device='cuda'): self.task = task self.device = device from dpt import DPTDepthModel if task == 'depth': path = 'pretrained/omnidata/omnidata_dpt_depth_v2.ckpt' self.model = DPTDepthModel(backbone='vitb_rn50_384') self.aug = transforms.Compose([ transforms.Resize((384, 384)), transforms.ToTensor(), transforms.Normalize(mean=0.5, std=0.5) ]) else: path = 'pretrained/omnidata/omnidata_dpt_normal_v2.ckpt' self.model = DPTDepthModel(backbone='vitb_rn50_384', num_channels=3) self.aug = transforms.Compose([ transforms.Resize((384, 384)), transforms.ToTensor() ]) checkpoint = torch.load(path, map_location='cpu') if 'state_dict' in checkpoint: state_dict = {} for k, v in checkpoint['state_dict'].items(): state_dict[k[6:]] = v else: state_dict = checkpoint self.model.load_state_dict(state_dict) self.model.eval().to(device) @torch.no_grad() def __call__(self, image): H, W = image.shape[:2] image = Image.fromarray(image) image = self.aug(image).unsqueeze(0).to(self.device) if self.task == 'depth': depth = self.model(image).clamp(0, 1) depth = F.interpolate(depth.unsqueeze(1), size=(H, W), mode='bicubic', align_corners=False) depth = depth.squeeze(1).cpu().numpy() return depth else: normal = self.model(image).clamp(0, 1) normal = F.interpolate(normal, size=(H, W), mode='bicubic', align_corners=False) normal = normal.cpu().numpy() return normal if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('path', type=str, help="path to image (png, jpeg, etc.)") parser.add_argument('--size', default=256, type=int, help="output resolution") parser.add_argument('--border_ratio', default=0.2, type=float, help="output border ratio") parser.add_argument('--recenter', type=bool, default=True, help="recenter, potentially not helpful for multiview zero123") parser.add_argument('--dont_recenter', dest='recenter', action='store_false') opt = parser.parse_args() out_dir = os.path.dirname(opt.path) out_rgba = os.path.join(out_dir, os.path.basename(opt.path).split('.')[0] + '_rgba.png') out_depth = os.path.join(out_dir, os.path.basename(opt.path).split('.')[0] + '_depth.png') out_normal = os.path.join(out_dir, os.path.basename(opt.path).split('.')[0] + '_normal.png') out_caption = os.path.join(out_dir, os.path.basename(opt.path).split('.')[0] + '_caption.txt') print(f'loading image...') image = cv2.imread(opt.path, cv2.IMREAD_UNCHANGED) if image.shape[-1] == 4: image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB) else: image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) print(f'[INFO] background removal...') carved_image = BackgroundRemoval()(image) # [H, W, 4] mask = carved_image[..., -1] > 0 print(f'[INFO] depth estimation...') dpt_depth_model = DPT(task='depth') depth = dpt_depth_model(image)[0] depth[mask] = (depth[mask] - depth[mask].min()) / (depth[mask].max() - depth[mask].min() + 1e-9) depth[~mask] = 0 depth = (depth * 255).astype(np.uint8) del dpt_depth_model print(f'[INFO] normal estimation...') dpt_normal_model = DPT(task='normal') normal = dpt_normal_model(image)[0] normal = (normal * 255).astype(np.uint8).transpose(1, 2, 0) normal[~mask] = 0 del dpt_normal_model if opt.recenter: print(f'[INFO] recenter...') final_rgba = np.zeros((opt.size, opt.size, 4), dtype=np.uint8) final_depth = np.zeros((opt.size, opt.size), dtype=np.uint8) final_normal = np.zeros((opt.size, opt.size, 3), dtype=np.uint8) coords = np.nonzero(mask) x_min, x_max = coords[0].min(), coords[0].max() y_min, y_max = coords[1].min(), coords[1].max() h = x_max - x_min w = y_max - y_min desired_size = int(opt.size * (1 - opt.border_ratio)) scale [truncated — 891 more characters] ``` ### stablediffusion/optimizer.py ```python import math from typing import List import torch from torch import Tensor from torch.optim.optimizer import Optimizer class Adan(Optimizer): def __init__(self, params, lr=1e-3, betas=(0.98, 0.92, 0.99), eps=1e-8, weight_decay=0.0, max_grad_norm=0.0, no_prox=False, foreach: bool = True): if not 0.0 <= max_grad_norm: raise ValueError('Invalid Max grad norm: {}'.format(max_grad_norm)) if not 0.0 <= lr: raise ValueError('Invalid learning rate: {}'.format(lr)) if not 0.0 <= eps: raise ValueError('Invalid epsilon value: {}'.format(eps)) if not 0.0 <= betas[0] < 1.0: raise ValueError('Invalid beta parameter at index 0: {}'.format( betas[0])) if not 0.0 <= betas[1] < 1.0: raise ValueError('Invalid beta parameter at index 1: {}'.format( betas[1])) if not 0.0 <= betas[2] < 1.0: raise ValueError('Invalid beta parameter at index 2: {}'.format( betas[2])) defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, max_grad_norm=max_grad_norm, no_prox=no_prox, foreach=foreach) super().__init__(params, defaults) def __setstate__(self, state): super(Adan, self).__setstate__(state) for group in self.param_groups: group.setdefault('no_prox', False) @torch.no_grad() def restart_opt(self): for group in self.param_groups: group['step'] = 0 for p in group['params']: if p.requires_grad: state = self.state[p] state['exp_avg'] = torch.zeros_like(p) state['exp_avg_sq'] = torch.zeros_like(p) state['exp_avg_diff'] = torch.zeros_like(p) @torch.no_grad() def step(self, closure=None): """Performs a single optimization step.""" loss = None if closure is not None: with torch.enable_grad(): loss = closure() if self.defaults['max_grad_norm'] > 0: device = self.param_groups[0]['params'][0].device global_grad_norm = torch.zeros(1, device=device) max_grad_norm = torch.tensor(self.defaults['max_grad_norm'], device=device) for group in self.param_groups: for p in group['params']: if p.grad is not None: grad = p.grad global_grad_norm.add_(grad.pow(2).sum()) global_grad_norm = torch.sqrt(global_grad_norm) clip_global_grad_norm = torch.clamp( max_grad_norm / (global_grad_norm + group['eps']), max=1.0).item() else: clip_global_grad_norm = 1.0 for group in self.param_groups: params_with_grad = [] grads = [] exp_avgs = [] exp_avg_sqs = [] exp_avg_diffs = [] neg_pre_grads = [] beta1, beta2, beta3 = group['betas'] if 'step' in group: group['step'] += 1 else: group['step'] = 1 bias_correction1 = 1.0 - beta1**group['step'] bias_correction2 = 1.0 - beta2**group['step'] bias_correction3 = 1.0 - beta3**group['step'] for p in group['params']: if p.grad is None: continue params_with_grad.append(p) grads.append(p.grad) state = self.state[p] if len(state) == 0: state['exp_avg'] = torch.zeros_like(p) state['exp_avg_sq'] = torch.zeros_like(p) state['exp_avg_diff'] = torch.zeros_like(p) if 'neg_pre_grad' not in state or group['step'] == 1: state['neg_pre_grad'] = p.grad.clone().mul_( -clip_global_grad_norm) exp_avgs.append(state['exp_avg']) exp_avg_sqs.append(state['exp_avg_sq']) exp_avg_diffs.append(state['exp_avg_diff']) neg_pre_grads.append(state['neg_pre_grad']) kwargs = dict( params=params_with_grad, grads=grads, exp_avgs=exp_avgs, exp_avg_sqs=exp_avg_sqs, exp_avg_diffs=exp_avg_diffs, neg_pre_grads=neg_pre_grads, beta1=beta1, beta2=beta2, beta3=beta3, bias_correction1=bias_correction1, bias_correction2=bias_correction2, bias_correction3_sqrt=math.sqrt(bias_correction3), lr=group['lr'], weight_decay=group['weight_decay'], eps=group['eps'], no_prox=group['no_prox'], clip_global_grad_norm=clip_global_grad_norm, ) if group['foreach']: _multi_tensor_adan(**kwargs) else: _single_tensor_adan(**kwargs) return loss def _single_tensor_adan( params: List[Tensor], grads: List[Tensor], exp_avgs: List[Tensor], exp_avg_sqs: List[Tensor], exp_avg_diffs: List[Tensor], neg_pre_grads: List[Tensor], *, beta1: float, beta2: float, beta3: float, bias_correction1: float, bias_correction2: float, bias_correction3_sqrt: float, lr: float, weight_decay: float, eps: float, no_prox: bool, clip_global_grad_norm: Tensor, ): for i, param in enumerate(params): grad = grads[i] exp_avg = exp_ [truncated — 3380 more characters] ``` ### stablediffusion/raymarching/__init__.py ```python from .raymarching import * ``` [38 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]