shap-e/shap_e/examples/gc/app.py

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)