Create app.py

app.py

@@ -0,0 +1,96 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import google.generativeai as genai
+from threading import Thread
+import trimesh
+import numpy as np
+import tempfile
+import os
+
+# Configure the API key securely from Streamlit's secrets
+genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
+
+# Set an environment variable for Hugging Face token
+os.environ["HF_TOKEN"] = st.secrets["HF_TOKEN"]
+
+# Load the LLaMA-Mesh model and tokenizer
+model_path = "Zhengyi/LLaMA-Mesh"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", low_cpu_mem_usage=True)
+terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
+
+def generate_mesh(prompt, temperature=0.9, max_new_tokens=4096):
+    conversation = [{"role": "user", "content": prompt}]
+    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        eos_token_id=terminators,
+    )
+
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+    return "".join(outputs)
+
+def apply_gradient_color(mesh_text):
+    temp_file = tempfile.NamedTemporaryFile(suffix="", delete=False).name
+    with open(temp_file + ".obj", "w") as f:
+        f.write(mesh_text)
+    mesh = trimesh.load_mesh(temp_file + ".obj", file_type='obj')
+
+    vertices = mesh.vertices
+    y_values = vertices[:, 1]
+
+    y_normalized = (y_values - y_values.min()) / (y_values.max() - y_values.min())
+
+    colors = np.zeros((len(vertices), 4))
+    colors[:, 0] = y_normalized
+    colors[:, 2] = 1 - y_normalized
+    colors[:, 3] = 1.0
+
+    mesh.visual.vertex_colors = colors
+
+    glb_path = temp_file + ".glb"
+    with open(glb_path, "wb") as f:
+        f.write(trimesh.exchange.gltf.export_glb(mesh))
+    return glb_path
+
+# Streamlit App UI
+st.title("Ever AI - 3D CAD Model Generator")
+st.write("Use generative AI to create 3D CAD models based on your prompt.")
+
+prompt = st.text_input("Enter your prompt:", "Create a 3D model of a house.")
+
+if st.button("Generate CAD Model"):
+    try:
+        response = generate_mesh(prompt)
+        
+        cad_file_path = "generated_model.obj"
+        with open(cad_file_path, "w") as f:
+            f.write(response)
+        
+        st.write("CAD Model Generated:")
+        st.code(response, language='plaintext')
+
+        glb_path = apply_gradient_color(response)
+        with open(glb_path, "rb") as f:
+            btn = st.download_button(
+                label="Download GLB File",
+                data=f,
+                file_name="generated_model.glb",
+                mime="application/octet-stream"
+            )
+    except Exception as e:
+        st.error(f"Error: {e}")