6 changed files with 508 additions and 15 deletions
@ -1,3 +1,5 @@ |
|||||
__pycache__/ |
__pycache__/ |
||||
.DS_Store |
.DS_Store |
||||
*.egg-info/ |
*.egg-info/ |
||||
|
shap_e/examples/shap_e_model_cache/ |
||||
|
shap_e/examples/gc/shap_e_model_cache/ |
||||
|
|||||
@ -0,0 +1,172 @@ |
|||||
|
{ |
||||
|
"cells": [ |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": 1, |
||||
|
"id": "964ccced", |
||||
|
"metadata": {}, |
||||
|
"outputs": [], |
||||
|
"source": [ |
||||
|
"import torch\n", |
||||
|
"\n", |
||||
|
"from shap_e.diffusion.sample import sample_latents\n", |
||||
|
"from shap_e.diffusion.gaussian_diffusion import diffusion_from_config\n", |
||||
|
"from shap_e.models.download import load_model, load_config\n", |
||||
|
"from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, gif_widget" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": 2, |
||||
|
"id": "8eed3a76", |
||||
|
"metadata": {}, |
||||
|
"outputs": [], |
||||
|
"source": [ |
||||
|
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": 3, |
||||
|
"id": "4127249f-da93-4da9-a15e-47fc1d918758", |
||||
|
"metadata": {}, |
||||
|
"outputs": [ |
||||
|
{ |
||||
|
"name": "stdout", |
||||
|
"output_type": "stream", |
||||
|
"text": [ |
||||
|
"NVIDIA GeForce RTX 3090\n" |
||||
|
] |
||||
|
} |
||||
|
], |
||||
|
"source": [ |
||||
|
"print(torch.cuda.get_device_name())" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": null, |
||||
|
"id": "2d922637", |
||||
|
"metadata": {}, |
||||
|
"outputs": [], |
||||
|
"source": [ |
||||
|
"xm = load_model('transmitter', device=device)\n", |
||||
|
"model = load_model('text300M', device=device)\n", |
||||
|
"diffusion = diffusion_from_config(load_config('diffusion'))" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": 19, |
||||
|
"id": "53d329d0", |
||||
|
"metadata": {}, |
||||
|
"outputs": [ |
||||
|
{ |
||||
|
"data": { |
||||
|
"application/vnd.jupyter.widget-view+json": { |
||||
|
"model_id": "f76a8f93c93e4b77af91f03645eb5011", |
||||
|
"version_major": 2, |
||||
|
"version_minor": 0 |
||||
|
}, |
||||
|
"text/plain": [ |
||||
|
" 0%| | 0/64 [00:00<?, ?it/s]" |
||||
|
] |
||||
|
}, |
||||
|
"metadata": {}, |
||||
|
"output_type": "display_data" |
||||
|
} |
||||
|
], |
||||
|
"source": [ |
||||
|
"batch_size = 1\n", |
||||
|
"guidance_scale = 30.0\n", |
||||
|
"prompt = \"bin bag\"\n", |
||||
|
"\n", |
||||
|
"latents = sample_latents(\n", |
||||
|
" batch_size=batch_size,\n", |
||||
|
" model=model,\n", |
||||
|
" diffusion=diffusion,\n", |
||||
|
" guidance_scale=guidance_scale,\n", |
||||
|
" model_kwargs=dict(texts=[prompt] * batch_size),\n", |
||||
|
" progress=True,\n", |
||||
|
" clip_denoised=True,\n", |
||||
|
" use_fp16=True,\n", |
||||
|
" use_karras=True,\n", |
||||
|
" karras_steps=64,\n", |
||||
|
" sigma_min=1e-3,\n", |
||||
|
" sigma_max=160,\n", |
||||
|
" s_churn=0,\n", |
||||
|
")" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": 20, |
||||
|
"id": "633da2ec", |
||||
|
"metadata": {}, |
||||
|
"outputs": [ |
||||
|
{ |
||||
|
"data": { |
||||
|
"application/vnd.jupyter.widget-view+json": { |
||||
|
"model_id": "7c8b8946a49847dd9aa5376f9568775f", |
||||
|
"version_major": 2, |
||||
|
"version_minor": 0 |
||||
|
}, |
||||
|
"text/plain": [ |
||||
|
"HTML(value='<img src=\"data:image/gif;base64,R0lGODlhQABAAIcAAJOYl4+UlY6TlI2TlI2Tk46Sko2Sk42SkoySk4ySko2RkYyRko…" |
||||
|
] |
||||
|
}, |
||||
|
"metadata": {}, |
||||
|
"output_type": "display_data" |
||||
|
} |
||||
|
], |
||||
|
"source": [ |
||||
|
"render_mode = 'nerf' # you can change this to 'stf'\n", |
||||
|
"size = 64 # this is the size of the renders; higher values take longer to render.\n", |
||||
|
"\n", |
||||
|
"cameras = create_pan_cameras(size, device)\n", |
||||
|
"for i, latent in enumerate(latents):\n", |
||||
|
" images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)\n", |
||||
|
" display(gif_widget(images))" |
||||
|
] |
||||
|
}, |
||||
|
{ |
||||
|
"cell_type": "code", |
||||
|
"execution_count": null, |
||||
|
"id": "85a4dce4", |
||||
|
"metadata": {}, |
||||
|
"outputs": [], |
||||
|
"source": [ |
||||
|
"# Example of saving the latents as meshes.\n", |
||||
|
"from shap_e.util.notebooks import decode_latent_mesh\n", |
||||
|
"\n", |
||||
|
"for i, latent in enumerate(latents):\n", |
||||
|
" t = decode_latent_mesh(xm, latent).tri_mesh()\n", |
||||
|
" with open(f'example_mesh_{i}.ply', 'wb') as f:\n", |
||||
|
" t.write_ply(f)\n", |
||||
|
" with open(f'example_mesh_{i}.obj', 'w') as f:\n", |
||||
|
" t.write_obj(f)" |
||||
|
] |
||||
|
} |
||||
|
], |
||||
|
"metadata": { |
||||
|
"kernelspec": { |
||||
|
"display_name": "Python 3 (ipykernel)", |
||||
|
"language": "python", |
||||
|
"name": "python3" |
||||
|
}, |
||||
|
"language_info": { |
||||
|
"codemirror_mode": { |
||||
|
"name": "ipython", |
||||
|
"version": 3 |
||||
|
}, |
||||
|
"file_extension": ".py", |
||||
|
"mimetype": "text/x-python", |
||||
|
"name": "python", |
||||
|
"nbconvert_exporter": "python", |
||||
|
"pygments_lexer": "ipython3", |
||||
|
"version": "3.10.12" |
||||
|
} |
||||
|
}, |
||||
|
"nbformat": 4, |
||||
|
"nbformat_minor": 5 |
||||
|
} |
||||
@ -0,0 +1,96 @@ |
|||||
|
import torch |
||||
|
from shap_e.diffusion.sample import sample_latents |
||||
|
from shap_e.diffusion.gaussian_diffusion import diffusion_from_config |
||||
|
from shap_e.models.download import load_model, load_config |
||||
|
from shap_e.util.notebooks import decode_latent_mesh |
||||
|
from tqdm import tqdm |
||||
|
import pygltflib |
||||
|
from pygltflib import GLTF2 |
||||
|
import trimesh |
||||
|
import open3d as o3d |
||||
|
import os |
||||
|
import datetime |
||||
|
|
||||
|
class ShapeGenerator: |
||||
|
def __init__(self, output_path, batch_size, step_size, guidance): |
||||
|
self.device = None |
||||
|
self.xm = None |
||||
|
self.model = None |
||||
|
self.diffusion = None |
||||
|
self.iterations = 0 |
||||
|
self.latents = None |
||||
|
self.output_path = output_path |
||||
|
self.batch_size = batch_size |
||||
|
self.step_size = step_size |
||||
|
self.guidance = guidance |
||||
|
|
||||
|
def run(self): |
||||
|
print("Loading Models..") |
||||
|
self.load_models() |
||||
|
print("Finished Loading Models!") |
||||
|
|
||||
|
def load_models(self): |
||||
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
||||
|
self.xm = load_model('transmitter', device=self.device) |
||||
|
self.model = load_model('text300M', device=self.device) |
||||
|
self.diffusion = diffusion_from_config(load_config('diffusion')) |
||||
|
|
||||
|
|
||||
|
def generate_object(self, prompt): |
||||
|
batch_size = 2 |
||||
|
# Create random latents |
||||
|
latent_dim = self.model.d_latent |
||||
|
random_latents = torch.randn(batch_size, latent_dim).to(self.model.device) |
||||
|
print(random_latents.shape) |
||||
|
model_kwargs = {} |
||||
|
|
||||
|
|
||||
|
self.latents = sample_latents( |
||||
|
batch_size=self.batch_size, |
||||
|
model=self.model, |
||||
|
diffusion=self.diffusion, |
||||
|
guidance_scale=self.guidance, |
||||
|
model_kwargs=model_kwargs, |
||||
|
progress=True, # This should already show progress |
||||
|
clip_denoised=True, |
||||
|
use_fp16=True, |
||||
|
use_karras=True, |
||||
|
karras_steps=self.step_size, |
||||
|
sigma_min=1e-3, |
||||
|
sigma_max=160, |
||||
|
s_churn=0, |
||||
|
device = self.model.device, |
||||
|
) |
||||
|
|
||||
|
self.export_model(prompt) |
||||
|
|
||||
|
def export_model(self, prompt): |
||||
|
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") |
||||
|
obj_filepath = f'{prompt}-{self.iterations}.obj' |
||||
|
output_filepath = f'{self.output_path}/{prompt}-{timestamp}.gltf' |
||||
|
print(output_filepath) |
||||
|
for i, latent in enumerate(self.latents): |
||||
|
t = decode_latent_mesh(self.xm, latent).tri_mesh() |
||||
|
with open(obj_filepath, 'w') as f: |
||||
|
t.write_obj(f) |
||||
|
final_mesh = self.construct_mesh(obj_filepath) |
||||
|
o3d.io.write_triangle_mesh(output_filepath, final_mesh) |
||||
|
self.iterations += 1 |
||||
|
|
||||
|
def construct_mesh(self, obj_fp): |
||||
|
mesh = o3d.io.read_triangle_mesh(obj_fp) |
||||
|
|
||||
|
if os.path.exists(obj_fp): |
||||
|
os.remove(obj_fp) |
||||
|
|
||||
|
original_triangle_count = len(mesh.triangles) |
||||
|
target_triangle_count = original_triangle_count // 3 |
||||
|
decimated_mesh = mesh.simplify_quadric_decimation( |
||||
|
target_number_of_triangles=target_triangle_count) |
||||
|
filtered_mesh = decimated_mesh.filter_smooth_simple(number_of_iterations=5) |
||||
|
filtered_mesh.compute_vertex_normals() |
||||
|
return filtered_mesh |
||||
|
|
||||
|
|
||||
|
|
||||
|
|
||||
@ -0,0 +1,51 @@ |
|||||
|
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, set_seed |
||||
|
import random |
||||
|
import torch |
||||
|
|
||||
|
class TextGenerator: |
||||
|
def __init__(self): |
||||
|
self.model = None |
||||
|
self.tokenizer = None |
||||
|
#self.load_models() |
||||
|
|
||||
|
def load_models(self): |
||||
|
print('Loading Models...') |
||||
|
self.tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") |
||||
|
self.model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") |
||||
|
print('Models Loaded!') |
||||
|
|
||||
|
def generate_text(self): |
||||
|
|
||||
|
model = AutoModelForCausalLM.from_pretrained( |
||||
|
"microsoft/Phi-3-mini-4k-instruct", |
||||
|
device_map="cuda", |
||||
|
torch_dtype="auto", |
||||
|
trust_remote_code=True, |
||||
|
) |
||||
|
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct") |
||||
|
|
||||
|
messages = [ |
||||
|
{"role": "system", "content": "You are a helpful AI assistant, that generates two nouns and returns one sentence in the format of: a (noun) with a (noun).\n You can descirbe a random object typically found in a bin"}, |
||||
|
{"role": "user", "content": "Can you provide me with a sentence"}, |
||||
|
] |
||||
|
|
||||
|
pipe = pipeline( |
||||
|
"text-generation", |
||||
|
model=model, |
||||
|
tokenizer=tokenizer, |
||||
|
) |
||||
|
|
||||
|
generation_args = { |
||||
|
"max_new_tokens": 50, # Reduced to focus on concise output |
||||
|
"return_full_text": False, |
||||
|
"temperature": 0.7, # Adjusted for more randomness |
||||
|
"do_sample": True, |
||||
|
"top_k": 100, # Top-k sampling |
||||
|
"top_p": 1, # Nucleus sampling |
||||
|
} |
||||
|
|
||||
|
output = pipe(messages, **generation_args) |
||||
|
return output[0]['generated_text'] |
||||
|
|
||||
|
|
||||
|
|
||||
@ -0,0 +1,105 @@ |
|||||
|
import argparse |
||||
|
import threading |
||||
|
import time |
||||
|
import random |
||||
|
from ShapeGenerator import ShapeGenerator |
||||
|
from TextGenerator import TextGenerator |
||||
|
|
||||
|
# command example |
||||
|
# python app.py --output_dir /mnt/c/Users/caile/Desktop/output |
||||
|
|
||||
|
class GCApp: |
||||
|
def __init__(self, output_dir, batch_size, step_size, guidance_scale): |
||||
|
self.output_dir = output_dir |
||||
|
self.obj_gen = ShapeGenerator(self.output_dir, batch_size, step_size, guidance_scale) |
||||
|
self.running = False |
||||
|
self.stop_event = threading.Event() |
||||
|
self.thread = None |
||||
|
self.waste_items = [ |
||||
|
"Plastic bottle", "Aluminum can", "Glass bottle", "Food wrapper", |
||||
|
"Cardboard box", "Paper bag", "Plastic bag", "Electronics", |
||||
|
"Old smartphone", "Broken TV", "Computer parts", "Batteries", |
||||
|
"Light bulbs", "Old furniture", "Styrofoam cup", "Food container", |
||||
|
"Takeout box", "Cigarette butts", "Plastic utensils", "Straws", |
||||
|
"Bottle caps", "Rubber tires", "Broken toys", "Old clothes", |
||||
|
"Shoes", "Wooden pallets", "Paint cans", "Cleaning products", |
||||
|
"Old appliances", "Wires", "Cables", "Extension cords", |
||||
|
"Old magazines", "Newspapers", "Scrap metal", "Construction debris", |
||||
|
"Yard waste", "Grass clippings", "Leaves", "Old mattresses", |
||||
|
"Carpeting", "Food scraps", "Pet waste", "Diapers", |
||||
|
"Sanitary products", "Receipts", "Plastic wrap", "Packing peanuts", |
||||
|
"Ice cream containers", "Fast food containers", "Takeaway cups", |
||||
|
"Clamshell packaging", "Plastic film", "Broken glass", |
||||
|
"Old books", "VCR tapes", "CDs", "DVDs", |
||||
|
"Game consoles", "Remote controls", "Ink cartridges", |
||||
|
"Toner cartridges", "Old tools", "Gardening tools", |
||||
|
"Bike parts", "Fishing gear", "Beach toys", "Pool floats", |
||||
|
"Old bicycles", "Skateboards", "Surfboards", "Helmets", |
||||
|
"Used batteries", "Old jewelry", "Keyboards", "Mice (computer)", |
||||
|
"Speakers", "Old cameras", "Projectors", "Printers", |
||||
|
"Scanners", "Shredded paper", "Bubble wrap", "Plastic sheeting", |
||||
|
"Tarps", "Old car parts", "Motor oil containers", |
||||
|
"Propane tanks", "Oil filters", "Windshield wipers", |
||||
|
"Car batteries", "Antifreeze containers", "Used tires", |
||||
|
"Old propane tanks", "Scrap wood", "Broken furniture", |
||||
|
"Old carpets", "Leather scraps", "Textile waste", |
||||
|
"Compostable waste" |
||||
|
] |
||||
|
|
||||
|
def start_generation(self): |
||||
|
self.running = True |
||||
|
self.stop_event.clear() |
||||
|
self.thread = threading.Thread(target=self._generate_objects) |
||||
|
self.thread.start() |
||||
|
|
||||
|
def stop_generation(self): |
||||
|
self.stop_event.set() |
||||
|
self.running = False |
||||
|
if self.thread: |
||||
|
self.thread.join() |
||||
|
|
||||
|
def get_random_item_prompt(self): |
||||
|
return random.choice(self.waste_items) |
||||
|
|
||||
|
def _generate_objects(self): |
||||
|
while not self.stop_event.is_set(): |
||||
|
self.obj_gen.generate_object(self.get_random_item_prompt()) |
||||
|
time.sleep(1) |
||||
|
|
||||
|
def run(self): |
||||
|
self.obj_gen.run() |
||||
|
|
||||
|
while True: |
||||
|
command = input("Enter a command, <start> <stop> <generate (prompt)>: ") |
||||
|
if command.lower() == 'exit': |
||||
|
print("Exiting the program.") |
||||
|
self.stop_generation() |
||||
|
break |
||||
|
elif command.lower() == 'start': |
||||
|
if not self.running: |
||||
|
print("Starting continuous generation.") |
||||
|
self.start_generation() |
||||
|
else: |
||||
|
print("Generation already running.") |
||||
|
elif command.lower() == 'stop': |
||||
|
print("Stopping continuous generation.") |
||||
|
self.stop_generation() |
||||
|
else: |
||||
|
print("Unknown command.") |
||||
|
|
||||
|
def main(output_dir, batch_size, step_size, guidance_scale): |
||||
|
app = GCApp(output_dir, batch_size, step_size, guidance_scale) |
||||
|
app.run() |
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
parser = argparse.ArgumentParser(description="Generate shapes with the ShapeGenerator.") |
||||
|
parser.add_argument("--output_dir", type=str, required=True, help="The directory to save generated shapes.") |
||||
|
parser.add_argument("--batch_size", type=int, default=2, help="The number of batches for shap-e. the higher the batch size the longer it will take to process but will output a more refined mesh.") |
||||
|
parser.add_argument("--step_size", type=int, default=64, help="The number of steps/iterations for shap-e. the higher the step size the longer it will take to process but will output a more refined mesh.") |
||||
|
parser.add_argument("--guidance_scale", type=int, default=30, help="The guidance scale in context to the text prompt. The higher this value, the model will generate something closer to the text description (CLIP).") |
||||
|
|
||||
|
args = parser.parse_args() |
||||
|
|
||||
|
main(args.output_dir, args.batch_size, args.step_size, args.guidance_scale) |
||||
|
|
||||
|
|
||||
Loading…
Reference in new issue