Update app.py

app.py

@@ -4,6 +4,7 @@ import uuid
 import json
 import time
 import asyncio
+import tempfile
 from threading import Thread
 
 import gradio as gr
@@ -12,6 +13,7 @@ import torch
 import numpy as np
 from PIL import Image
 import edge_tts
+import trimesh
 
 from transformers import (
     AutoModelForCausalLM,
@@ -21,14 +23,75 @@ from transformers import (
     AutoProcessor,
 )
 from transformers.image_utils import load_image
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-# Additional imports for 3D model generation
-import tempfile
-import trimesh
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from diffusers import ShapEImg2ImgPipeline, ShapEPipeline
 from diffusers.utils import export_to_ply
 
+# -----------------------------------------------------------------------------
+# Global constants and helper functions
+# -----------------------------------------------------------------------------
+
+MAX_SEED = np.iinfo(np.int32).max
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+# -----------------------------------------------------------------------------
+# Model class for Text-to-3D Generation (ShapE)
+# -----------------------------------------------------------------------------
+
+class Model:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16)
+        self.pipe.to(self.device)
+
+        self.pipe_img = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
+        self.pipe_img.to(self.device)
+
+    def to_glb(self, ply_path: str) -> str:
+        mesh = trimesh.load(ply_path)
+        rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
+        mesh.apply_transform(rot)
+        rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
+        mesh.apply_transform(rot)
+        mesh_path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
+        mesh.export(mesh_path.name, file_type="glb")
+        return mesh_path.name
+
+    def run_text(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64) -> str:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        images = self.pipe(
+            prompt,
+            generator=generator,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+            output_type="mesh",
+        ).images
+        ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
+        export_to_ply(images[0], ply_path.name)
+        return self.to_glb(ply_path.name)
+
+    def run_image(self, image: Image.Image, seed: int = 0, guidance_scale: float = 3.0, num_steps: int = 64) -> str:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        images = self.pipe_img(
+            image,
+            generator=generator,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+            output_type="mesh",
+        ).images
+        ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
+        export_to_ply(images[0], ply_path.name)
+        return self.to_glb(ply_path.name)
+
+# -----------------------------------------------------------------------------
+# Gradio UI configuration
+# -----------------------------------------------------------------------------
+
 DESCRIPTION = """
 # QwQ Edge 💬
 """
@@ -50,13 +113,14 @@ h1 {
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-MAX_SEED = np.iinfo(np.int32).max
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# ------------------------------
-# Text Generation Model
-# ------------------------------
+# -----------------------------------------------------------------------------
+# Load Models and Pipelines for Chat, Image, and Multimodal Processing
+# -----------------------------------------------------------------------------
+
+# Load the text-only model and tokenizer (for pure text chat)
 model_id = "prithivMLmods/FastThink-0.5B-Tiny"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
@@ -66,11 +130,13 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+# Voices for text-to-speech
 TTS_VOICES = [
     "en-US-JennyNeural",  # @tts1
     "en-US-GuyNeural",    # @tts2
 ]
 
+# Load multimodal processor and model (e.g. for OCR and image processing)
 MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 model_m = Qwen2VLForConditionalGeneration.from_pretrained(
@@ -79,12 +145,20 @@ model_m = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to("cuda").eval()
 
+# -----------------------------------------------------------------------------
+# Asynchronous text-to-speech
+# -----------------------------------------------------------------------------
+
 async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
     """Convert text to speech using Edge TTS and save as MP3"""
     communicate = edge_tts.Communicate(text, voice)
     await communicate.save(output_file)
     return output_file
 
+# -----------------------------------------------------------------------------
+# Utility function to clean conversation history
+# -----------------------------------------------------------------------------
+
 def clean_chat_history(chat_history):
     """
     Filter out any chat entries whose "content" is not a string.
@@ -96,9 +170,10 @@ def clean_chat_history(chat_history):
             cleaned.append(msg)
     return cleaned
 
-# ------------------------------
-# Stable Diffusion XL (Image Generation)
-# ------------------------------
+# -----------------------------------------------------------------------------
+# Stable Diffusion XL Pipeline for Image Generation
+# -----------------------------------------------------------------------------
+
 MODEL_ID_SD = os.getenv("MODEL_VAL_PATH")  # SDXL Model repository path via env variable
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
@@ -128,11 +203,6 @@ def save_image(img: Image.Image) -> str:
     img.save(unique_name)
     return unique_name
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
 @spaces.GPU(duration=60, enable_queue=True)
 def generate_image_fn(
     prompt: str,
@@ -166,6 +236,7 @@ def generate_image_fn(
         options["use_resolution_binning"] = True
 
     images = []
+    # Process in batches
     for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
@@ -180,82 +251,30 @@ def generate_image_fn(
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-# ------------------------------
-# 3D Model Generation using ShapE (Text-to-3D / Image-to-3D)
-# ------------------------------
-class Model3D:
-    def __init__(self):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16)
-        self.pipe.to(self.device)
-        # Ensure the text encoder is in half precision
-        self.pipe.text_encoder = self.pipe.text_encoder.half()
-
-        self.pipe_img = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
-        self.pipe_img.to(self.device)
-        # Ensure the text encoder is in half precision
-        self.pipe_img.text_encoder = self.pipe_img.text_encoder.half()
-
-    def to_glb(self, ply_path: str) -> str:
-        mesh = trimesh.load(ply_path)
-        rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
-        mesh.apply_transform(rot)
-        rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
-        mesh.apply_transform(rot)
-        mesh_path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
-        mesh.export(mesh_path.name, file_type="glb")
-        return mesh_path.name
+# -----------------------------------------------------------------------------
+# Text-to-3D Generation using the ShapE Pipeline
+# -----------------------------------------------------------------------------
 
-    def run_text(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64) -> str:
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-        if self.device.type == "cuda":
-            with torch.autocast("cuda", dtype=torch.float16):
-                output = self.pipe(
-                    prompt,
-                    generator=generator,
-                    guidance_scale=guidance_scale,
-                    num_inference_steps=num_steps,
-                    output_type="mesh",
-                )
-        else:
-            output = self.pipe(
-                prompt,
-                generator=generator,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_steps,
-                output_type="mesh",
-            )
-        images = output.images
-        ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
-        export_to_ply(images[0], ply_path.name)
-        return self.to_glb(ply_path.name)
-
-    def run_image(self, image: Image.Image, seed: int = 0, guidance_scale: float = 3.0, num_steps: int = 64) -> str:
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-        if self.device.type == "cuda":
-            with torch.autocast("cuda", dtype=torch.float16):
-                output = self.pipe_img(
-                    image,
-                    generator=generator,
-                    guidance_scale=guidance_scale,
-                    num_inference_steps=num_steps,
-                    output_type="mesh",
-                )
-        else:
-            output = self.pipe_img(
-                image,
-                generator=generator,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_steps,
-                output_type="mesh",
-            )
-        images = output.images
-        ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
-        export_to_ply(images[0], ply_path.name)
-        return self.to_glb(ply_path.name)
+@spaces.GPU(duration=120, enable_queue=True)
+def generate_3d_fn(
+    prompt: str,
+    seed: int = 1,
+    guidance_scale: float = 15.0,
+    num_steps: int = 64,
+    randomize_seed: bool = False,
+):
+    """
+    Generate a 3D model from text using the ShapE pipeline.
+    Returns a tuple of (glb_file_path, used_seed).
+    """
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    model3d = Model()
+    glb_path = model3d.run_text(prompt, seed=seed, guidance_scale=guidance_scale, num_steps=num_steps)
+    return glb_path, seed
 
-# Create a global instance of the 3D model generator.
-model_3d = Model3D()
+# -----------------------------------------------------------------------------
+# Chat Generation Function with support for @tts, @image, and @3d commands
+# -----------------------------------------------------------------------------
 
 @spaces.GPU
 def generate(
@@ -269,7 +288,7 @@ def generate(
 ):
     """
     Generates chatbot responses with support for multimodal input, TTS, image generation,
-    and 3D model generation.
+    and now 3D generation.
     
     Special commands:
       - "@tts1" or "@tts2": triggers text-to-speech.
@@ -279,24 +298,21 @@ def generate(
     text = input_dict["text"]
     files = input_dict.get("files", [])
 
-    # ------------------------------
-    # 3D Model Generation Command
-    # ------------------------------
+    # --- 3D Generation branch ---
     if text.strip().lower().startswith("@3d"):
-        text = text[len("@3d"):].strip()
+        prompt = text[len("@3d"):].strip()
         yield "Generating 3D model..."
-        seed = random.randint(0, MAX_SEED)
-        if files:
-            image = load_image(files[0])
-            glb_file = model_3d.run_image(image, seed=seed)
-        else:
-            glb_file = model_3d.run_text(text, seed=seed)
-        yield gr.File(glb_file)
+        glb_path, used_seed = generate_3d_fn(
+            prompt=prompt,
+            seed=1,
+            guidance_scale=15.0,
+            num_steps=64,
+            randomize_seed=True,
+        )
+        yield gr.File(glb_path, label="3D Model (GLB)")
         return
 
-    # ------------------------------
-    # Image Generation Command
-    # ------------------------------
+    # --- Image Generation branch ---
     if text.strip().lower().startswith("@image"):
         prompt = text[len("@image"):].strip()
         yield "Generating image..."
@@ -316,9 +332,7 @@ def generate(
         yield gr.Image(image_paths[0])
         return
 
-    # ------------------------------
-    # TTS / Regular Text Generation
-    # ------------------------------
+    # --- Text and TTS branch ---
     tts_prefix = "@tts"
     is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 3))
     voice_index = next((i for i in range(1, 3) if text.strip().lower().startswith(f"{tts_prefix}{i}")), None)
@@ -394,6 +408,10 @@ def generate(
             output_file = asyncio.run(text_to_speech(final_response, voice))
             yield gr.Audio(output_file, autoplay=True)
 
+# -----------------------------------------------------------------------------
+# Gradio Chat Interface Setup and Launch
+# -----------------------------------------------------------------------------
+
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
@@ -408,9 +426,9 @@ demo = gr.ChatInterface(
         [{"text": "Extract JSON from the image", "files": ["examples/document.jpg"]}],
         [{"text": "summarize the letter", "files": ["examples/1.png"]}],
         ["@image Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic"],
+        ["@3d A futuristic city skyline in the style of cyberpunk"],
         ["Write a Python function to check if a number is prime."],
         ["@tts2 What causes rainbows to form?"],
-        ["@3d A futuristic spaceship in low-poly style"],
     ],
     cache_examples=False,
     type="messages",
@@ -423,4 +441,5 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
+    # To create a public link, set share=True in launch().
     demo.queue(max_size=20).launch(share=True)