Update app.py

app.py

@@ -4,7 +4,6 @@ import uuid
 import json
 import time
 import asyncio
-import tempfile
 from threading import Thread
 
 import gradio as gr
@@ -13,7 +12,6 @@ import torch
 import numpy as np
 from PIL import Image
 import edge_tts
-import trimesh
 
 from transformers import (
     AutoModelForCausalLM,
@@ -23,75 +21,8 @@ from transformers import (
     AutoProcessor,
 )
 from transformers.image_utils import load_image
-
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from diffusers import ShapEImg2ImgPipeline, ShapEPipeline
-from diffusers.utils import export_to_ply
-
-
-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
 
-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)
-        # Ensure the text encoder is in half precision to avoid dtype mismatches.
-        if torch.cuda.is_available():
-            try:
-                self.pipe.text_encoder = self.pipe.text_encoder.half()
-            except AttributeError:
-                pass
-
-        self.pipe_img = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
-        self.pipe_img.to(self.device)
-        # Use getattr with a default value to avoid AttributeError if text_encoder is missing.
-        if torch.cuda.is_available():
-            text_encoder_img = getattr(self.pipe_img, "text_encoder", None)
-            if text_encoder_img is not None:
-                self.pipe_img.text_encoder = text_encoder_img.half()
-
-    def to_glb(self, ply_path: str) -> str:
-        mesh = trimesh.load(ply_path)
-        # Rotate the mesh for proper orientation
-        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)
 
 DESCRIPTION = """
 # QwQ Edge 💬
@@ -117,7 +48,7 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Load the text-only model and tokenizer (for pure text chat)
+# Load text-only model and tokenizer
 model_id = "prithivMLmods/FastThink-0.5B-Tiny"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
@@ -127,13 +58,11 @@ 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(
@@ -159,12 +88,14 @@ def clean_chat_history(chat_history):
             cleaned.append(msg)
     return cleaned
 
+# Environment variables and parameters for Stable Diffusion XL
 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"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  # For batched image generation
 
+# Load the SDXL pipeline
 sd_pipe = StableDiffusionXLPipeline.from_pretrained(
     MODEL_ID_SD,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -173,21 +104,31 @@ sd_pipe = StableDiffusionXLPipeline.from_pretrained(
 ).to(device)
 sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config)
 
+# Ensure that the text encoder is in half-precision if using CUDA.
 if torch.cuda.is_available():
     sd_pipe.text_encoder = sd_pipe.text_encoder.half()
 
+# Optional: compile the model for speedup if enabled
 if USE_TORCH_COMPILE:
     sd_pipe.compile()
 
+# Optional: offload parts of the model to CPU if needed
 if ENABLE_CPU_OFFLOAD:
     sd_pipe.enable_model_cpu_offload()
 
+MAX_SEED = np.iinfo(np.int32).max
+
 def save_image(img: Image.Image) -> str:
     """Save a PIL image with a unique filename and return the path."""
     unique_name = str(uuid.uuid4()) + ".png"
     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,
@@ -227,6 +168,7 @@ def generate_image_fn(
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options and batch_options["negative_prompt"] is not None:
             batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
+        # Wrap the pipeline call in autocast if using CUDA
         if device.type == "cuda":
             with torch.autocast("cuda", dtype=torch.float16):
                 outputs = sd_pipe(**batch_options)
@@ -236,23 +178,6 @@ def generate_image_fn(
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-@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
-
 @spaces.GPU
 def generate(
     input_dict: dict,
@@ -264,34 +189,16 @@ def generate(
     repetition_penalty: float = 1.2,
 ):
     """
-    Generates chatbot responses with support for multimodal input, TTS, image generation,
-    and 3D model generation.
-    
+    Generates chatbot responses with support for multimodal input, TTS, and image generation.
     Special commands:
       - "@tts1" or "@tts2": triggers text-to-speech.
       - "@image": triggers image generation using the SDXL pipeline.
-      - "@3d": triggers 3D model generation using the ShapE pipeline.
     """
     text = input_dict["text"]
     files = input_dict.get("files", [])
 
-    # --- 3D Generation branch ---
-    if text.strip().lower().startswith("@3d"):
-        prompt = text[len("@3d"):].strip()
-        yield "Generating 3D model..."
-        glb_path, used_seed = generate_3d_fn(
-            prompt=prompt,
-            seed=1,
-            guidance_scale=15.0,
-            num_steps=64,
-            randomize_seed=True,
-        )
-        # Instead of returning as a file, yield a 3D model component so it displays inline.
-        yield gr.Model3D(value=glb_path, label="3D Model")
-        return
-
-    # --- Image Generation branch ---
     if text.strip().lower().startswith("@image"):
+        # Remove the "@image" tag and use the rest as prompt
         prompt = text[len("@image"):].strip()
         yield "Generating image..."
         image_paths, used_seed = generate_image_fn(
@@ -307,10 +214,10 @@ def generate(
             use_resolution_binning=True,
             num_images=1,
         )
+        # Yield the generated image so that the chat interface displays it.
         yield gr.Image(image_paths[0])
-        return
+        return  # Exit early
 
-    # --- 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)
@@ -318,9 +225,11 @@ def generate(
     if is_tts and voice_index:
         voice = TTS_VOICES[voice_index - 1]
         text = text.replace(f"{tts_prefix}{voice_index}", "").strip()
+        # Clear previous chat history for a fresh TTS request.
         conversation = [{"role": "user", "content": text}]
     else:
         voice = None
+        # Remove any stray @tts tags and build the conversation history.
         text = text.replace(tts_prefix, "").strip()
         conversation = clean_chat_history(chat_history)
         conversation.append({"role": "user", "content": text})
@@ -354,6 +263,7 @@ def generate(
             time.sleep(0.01)
             yield buffer
     else:
+        
         input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
         if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -382,6 +292,7 @@ def generate(
         final_response = "".join(outputs)
         yield final_response
 
+        # If TTS was requested, convert the final response to speech.
         if is_tts and voice:
             output_file = asyncio.run(text_to_speech(final_response, voice))
             yield gr.Audio(output_file, autoplay=True)
@@ -397,11 +308,12 @@ demo = gr.ChatInterface(
     ],
     examples=[
         ["@tts1 Who is Nikola Tesla, and why did he die?"],
-        ["@3d A birthday cupcake with cherry"],
+        [{"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"],
         ["Write a Python function to check if a number is prime."],
         ["@tts2 What causes rainbows to form?"],
+
     ],
     cache_examples=False,
     type="messages",