Upload multimodal_module.py

multimodal_module.py

@@ -0,0 +1,644 @@
+# multimodal_module.py
+import os
+import pickle
+import subprocess
+import tempfile
+import shutil
+import asyncio
+from datetime import datetime
+from typing import Dict, List, Optional, Any
+import io
+import uuid
+
+# Core ML libs
+import torch
+from transformers import (
+    pipeline,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    Wav2Vec2Processor,
+    Wav2Vec2ForSequenceClassification,
+)
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer as HFTokenizer
+
+# Audio / speech
+import librosa
+import speech_recognition as sr
+from gtts import gTTS
+
+# Image, video, files
+from PIL import Image, ImageOps
+import imageio_ffmpeg as ffmpeg
+import imageio
+import moviepy.editor as mp
+import fitz  # PyMuPDF for PDFs
+
+# Misc
+from langdetect import DetectorFactory
+DetectorFactory.seed = 0
+
+# Optional: safety-check toggles
+USE_SAFETY_CHECKER = False
+
+# Helper for temp files
+def _tmp_path(suffix=""):
+    return os.path.join(tempfile.gettempdir(), f"{uuid.uuid4().hex}{suffix}")
+
+class MultiModalChatModule:
+    """
+    Full-power multimodal module.
+    - Lazy-loads big models on first use.
+    - Methods are async-friendly.
+    """
+
+    def __init__(self, chat_history_file: str = "chat_histories.pkl"):
+        self.user_chat_histories: Dict[int, List[dict]] = self._load_chat_histories(chat_history_file)
+        self.chat_history_file = chat_history_file
+
+        # device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"[MultiModal] device: {self.device}")
+
+        # placeholders for large models (lazy)
+        self._voice_processor = None
+        self._voice_emotion_model = None
+
+        self._translator = None
+
+        self._chat_tokenizer = None
+        self._chat_model = None
+        self._chat_model_name = "bigscience/bloom"  # placeholder; will set proper below
+
+        self._image_captioner = None
+
+        self._sd_pipe = None
+        self._sd_inpaint = None
+
+        self._code_tokenizer = None
+        self._code_model = None
+
+        # other small helpers
+        self._sr_recognizer = sr.Recognizer()
+
+        # set common model names (you can change)
+        self.model_names = {
+            "voice_emotion_processor": "facebook/hubert-large-ls960-ft",
+            "voice_emotion_model": "superb/hubert-base-superb-er",
+            "translation_model": "facebook/nllb-200-distilled-600M",
+            "chatbot_tokenizer": "facebook/blenderbot-400M-distill",
+            "chatbot_model": "facebook/blenderbot-400M-distill",
+            "image_captioner": "Salesforce/blip-image-captioning-base",
+            "sd_inpaint": "runwayml/stable-diffusion-inpainting",
+            "sd_text2img": "runwayml/stable-diffusion-v1-5",
+            "code_model": "bigcode/starcoder",  # Or use a specific StarCoder checkpoint on HF
+        }
+
+        # keep track of which heavy groups are loaded
+        self._loaded = {
+            "voice": False,
+            "translation": False,
+            "chat": False,
+            "image_caption": False,
+            "sd": False,
+            "code": False,
+        }
+
+    # ----------------------
+    # persistence
+    # ----------------------
+    def _load_chat_histories(self, fn: str) -> Dict[int, List[dict]]:
+        try:
+            with open(fn, "rb") as f:
+                return pickle.load(f)
+        except Exception:
+            return {}
+
+    def _save_chat_histories(self):
+        try:
+            with open(self.chat_history_file, "wb") as f:
+                pickle.dump(self.user_chat_histories, f)
+        except Exception as e:
+            print("[MultiModal] Warning: failed to save chat histories:", e)
+
+    # ----------------------
+    # Lazy loaders
+    # ----------------------
+    def _load_voice_models(self):
+        if self._loaded["voice"]:
+            return
+        print("[MultiModal] Loading voice/emotion models...")
+        self._voice_processor = Wav2Vec2Processor.from_pretrained(self.model_names["voice_emotion_processor"])
+        self._voice_emotion_model = Wav2Vec2ForSequenceClassification.from_pretrained(self.model_names["voice_emotion_model"]).to(self.device)
+        self._loaded["voice"] = True
+        print("[MultiModal] Voice models loaded.")
+
+    def _load_translation(self):
+        if self._loaded["translation"]:
+            return
+        print("[MultiModal] Loading translation pipeline...")
+        device_idx = 0 if self.device == "cuda" else -1
+        self._translator = pipeline("translation", model=self.model_names["translation_model"], device=device_idx)
+        self._loaded["translation"] = True
+        print("[MultiModal] Translation loaded.")
+
+    def _load_chatbot(self):
+        if self._loaded["chat"]:
+            return
+        print("[MultiModal] Loading chatbot model...")
+        # chatbot: keep current blenderbot to preserve behaviour
+        self._chat_tokenizer = AutoTokenizer.from_pretrained(self.model_names["chatbot_tokenizer"])
+        self._chat_model = AutoModelForSeq2SeqLM.from_pretrained(self.model_names["chatbot_model"]).to(self.device)
+        self._loaded["chat"] = True
+        print("[MultiModal] Chatbot loaded.")
+
+    def _load_image_captioner(self):
+        if self._loaded["image_caption"]:
+            return
+        print("[MultiModal] Loading image captioner...")
+        device_idx = 0 if self.device == "cuda" else -1
+        self._image_captioner = pipeline("image-to-text", model=self.model_names["image_captioner"], device=device_idx)
+        self._loaded["image_caption"] = True
+        print("[MultiModal] Image captioner loaded.")
+
+    def _load_sd(self):
+        if self._loaded["sd"]:
+            return
+        print("[MultiModal] Loading Stable Diffusion pipelines...")
+        # text2img
+        sd_model = self.model_names["sd_text2img"]
+        sd_inpaint_model = self.model_names["sd_inpaint"]
+        # Use float16 on GPU for speed
+        torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
+        try:
+            self._sd_pipe = StableDiffusionPipeline.from_pretrained(sd_model, torch_dtype=torch_dtype)
+            self._sd_pipe = self._sd_pipe.to(self.device)
+        except Exception as e:
+            print("[MultiModal] Warning loading text2img:", e)
+            self._sd_pipe = None
+
+        try:
+            self._sd_inpaint = StableDiffusionInpaintPipeline.from_pretrained(sd_inpaint_model, torch_dtype=torch_dtype)
+            self._sd_inpaint = self._sd_inpaint.to(self.device)
+        except Exception as e:
+            print("[MultiModal] Warning loading inpaint:", e)
+            self._sd_inpaint = None
+
+        self._loaded["sd"] = True
+        print("[MultiModal] Stable Diffusion loaded (where possible).")
+
+    def _load_code_model(self):
+        if self._loaded["code"]:
+            return
+        print("[MultiModal] Loading code model...")
+        # StarCoder style model (may require HF_TOKEN or large memory)
+        try:
+            self._code_tokenizer = HFTokenizer.from_pretrained(self.model_names["code_model"])
+            self._code_model = AutoModelForCausalLM.from_pretrained(self.model_names["code_model"]).to(self.device)
+            self._loaded["code"] = True
+            print("[MultiModal] Code model loaded.")
+        except Exception as e:
+            print("[MultiModal] Warning: could not load code model:", e)
+            self._code_tokenizer = None
+            self._code_model = None
+
+    # ----------------------
+    # Voice: analyze emotion, transcribe
+    # ----------------------
+    async def analyze_voice_emotion(self, audio_path: str) -> str:
+        self._load_voice_models()
+        speech, sr_ = librosa.load(audio_path, sr=16000)
+        inputs = self._voice_processor(speech, sampling_rate=sr_, return_tensors="pt", padding=True).to(self.device)
+        with torch.no_grad():
+            logits = self._voice_emotion_model(**inputs).logits
+        predicted_class = torch.argmax(logits).item()
+        return {
+            0: "😊 Happy",
+            1: "😢 Sad",
+            2: "😠 Angry",
+            3: "😨 Fearful",
+            4: "😌 Calm",
+            5: "😲 Surprised",
+        }.get(predicted_class, "🤔 Unknown")
+
+    async def process_voice_message(self, voice_file, user_id: int) -> dict:
+        """
+        voice_file: Starlette UploadFile or object with get_file() used previously in your code.
+        Returns: {text, language, emotion}
+        """
+        # Save OGG locally
+        ogg_path = _tmp_path(".ogg")
+        wav_path = _tmp_path(".wav")
+        tf = await voice_file.get_file()
+        await tf.download_to_drive(ogg_path)
+
+        # Convert to WAV via ffmpeg
+        try:
+            ffmpeg_path = ffmpeg.get_ffmpeg_exe()
+            subprocess.run([ffmpeg_path, "-y", "-i", ogg_path, wav_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        except Exception as e:
+            # fallback: try ffmpeg in PATH
+            try:
+                subprocess.run(["ffmpeg", "-y", "-i", ogg_path, wav_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            except Exception as ee:
+                raise RuntimeError(f"ffmpeg conversion failed: {e} / {ee}")
+
+        # Transcribe using SpeechRecognition Google STT (as before) -- or you can integrate whisper
+        recognizer = self._sr_recognizer
+        with sr.AudioFile(wav_path) as source:
+            audio = recognizer.record(source)
+
+        detected_lang = None
+        detected_text = ""
+        # tried languages set
+        lang_map = {
+            "zh": {"stt": "zh-CN"},
+            "ja": {"stt": "ja-JP"},
+            "ko": {"stt": "ko-KR"},
+            "en": {"stt": "en-US"},
+            "es": {"stt": "es-ES"},
+            "fr": {"stt": "fr-FR"},
+            "de": {"stt": "de-DE"},
+            "it": {"stt": "it-IT"},
+        }
+        for lang_code, lang_data in lang_map.items():
+            try:
+                detected_text = recognizer.recognize_google(audio, language=lang_data["stt"])
+                detected_lang = lang_code
+                break
+            except sr.UnknownValueError:
+                continue
+            except Exception:
+                continue
+
+        if not detected_lang:
+            # If not recognized, try fallback: detect from small chunk via langdetect
+            detected_lang = "en"
+            detected_text = ""
+
+        # emotion
+        emotion = await self.analyze_voice_emotion(wav_path)
+
+        # remove temp files
+        try:
+            os.remove(ogg_path)
+            os.remove(wav_path)
+        except Exception:
+            pass
+
+        return {"text": detected_text, "language": detected_lang, "emotion": emotion}
+
+    # ----------------------
+    # Text chat with translation & history
+    # ----------------------
+    async def generate_response(self, text: str, user_id: int, lang: str = "en") -> str:
+        # Ensure chat model loaded
+        self._load_chatbot()
+        self._load_translation()
+
+        if user_id not in self.user_chat_histories:
+            self.user_chat_histories[user_id] = []
+
+        self.user_chat_histories[user_id].append({"timestamp": datetime.now().isoformat(), "role": "user", "text": text, "language": lang})
+        self.user_chat_histories[user_id] = self.user_chat_histories[user_id][-100:]
+        self._save_chat_histories()
+
+        # Build context: translate last few msgs to English for consistency
+        context_texts = []
+        for msg in self.user_chat_histories[user_id][-5:]:
+            if msg.get("language", "en") != "en":
+                try:
+                    translated = self._translator(msg["text"])[0]["translation_text"]
+                except Exception:
+                    translated = msg["text"]
+            else:
+                translated = msg["text"]
+            context_texts.append(f"{msg['role']}: {translated}")
+
+        context = "\n".join(context_texts)
+        input_text = f"Context:\n{context}\nUser: {text if lang == 'en' else context_texts[-1].split(': ', 1)[1]}"
+
+        # Tokenize + generate
+        inputs = self._chat_tokenizer.encode(input_text, return_tensors="pt").to(self.device)
+        outputs = self._chat_model.generate(inputs, max_length=1000)
+        response_en = self._chat_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Translate back to user's language if needed
+        if lang != "en":
+            try:
+                response = self._translator(response_en)[0]["translation_text"]
+            except Exception:
+                response = response_en
+        else:
+            response = response_en
+
+        self.user_chat_histories[user_id].append({"timestamp": datetime.now().isoformat(), "role": "bot", "text": response, "language": lang})
+        self._save_chat_histories()
+
+        return response
+
+    # ----------------------
+    # Image captioning (existing)
+    # ----------------------
+    async def process_image_message(self, image_file, user_id: int) -> str:
+        # Save image
+        img_path = _tmp_path(".jpg")
+        tf = await image_file.get_file()
+        await tf.download_to_drive(img_path)
+
+        # load captioner
+        self._load_image_captioner()
+        try:
+            image = Image.open(img_path).convert("RGB")
+            description = self._image_captioner(image)[0]["generated_text"]
+        except Exception as e:
+            description = f"[Error generating caption: {e}]"
+
+        # cleanup
+        try:
+            os.remove(img_path)
+        except Exception:
+            pass
+
+        # store in history
+        if user_id not in self.user_chat_histories:
+            self.user_chat_histories[user_id] = []
+        self.user_chat_histories[user_id].append({"timestamp": datetime.now().isoformat(), "role": "user", "text": "[Image]", "language": "en"})
+        self.user_chat_histories[user_id].append({"timestamp": datetime.now().isoformat(), "role": "bot", "text": f"Image description: {description}", "language": "en"})
+        self._save_chat_histories()
+
+        return description
+
+    # ----------------------
+    # Voice reply (TTS)
+    # ----------------------
+    async def generate_voice_reply(self, text: str, user_id: int, fmt: str = "ogg") -> str:
+        """
+        Generate TTS audio reply using gTTS (or swap out to another TTS if you have).
+        Returns path to audio file.
+        """
+        mp3_path = _tmp_path(".mp3")
+        out_path = _tmp_path(f".{fmt}")
+
+        try:
+            tts = gTTS(text)
+            tts.save(mp3_path)
+            # convert to requested format using ffmpeg (ogg/opus for Telegram voice)
+            ffmpeg_path = ffmpeg.get_ffmpeg_exe()
+            if fmt == "ogg":
+                # convert mp3 -> ogg (opus)
+                subprocess.run([ffmpeg_path, "-y", "-i", mp3_path, "-c:a", "libopus", out_path], check=True)
+            elif fmt == "wav":
+                subprocess.run([ffmpeg_path, "-y", "-i", mp3_path, out_path], check=True)
+            else:
+                # default: return mp3
+                shutil.move(mp3_path, out_path)
+        except Exception as e:
+            # fallback: raise
+            raise RuntimeError(f"TTS failed: {e}")
+        finally:
+            try:
+                if os.path.exists(mp3_path) and os.path.exists(out_path) and mp3_path != out_path:
+                    os.remove(mp3_path)
+            except Exception:
+                pass
+
+        return out_path
+
+    # ----------------------
+    # Image generation (text -> image)
+    # ----------------------
+    async def generate_image_from_text(self, prompt: str, user_id: int, width: int = 512, height: int = 512, steps: int = 30) -> str:
+        self._load_sd()
+        if self._sd_pipe is None:
+            raise RuntimeError("Stable Diffusion pipeline not available.")
+
+        out_path = _tmp_path(".png")
+        try:
+            # diffusion pipeline uses CPU/GPU internally
+            result = self._sd_pipe(prompt, num_inference_steps=steps, height=height, width=width)
+            image = result.images[0]
+            image.save(out_path)
+        except Exception as e:
+            raise RuntimeError(f"Image generation failed: {e}")
+
+        return out_path
+
+    # ----------------------
+    # Image editing (inpainting)
+    # ----------------------
+    async def edit_image_inpaint(self, image_file, mask_file=None, prompt: str = "", user_id: int = 0) -> str:
+        self._load_sd()
+        if self._sd_inpaint is None:
+            raise RuntimeError("Inpainting pipeline not available.")
+
+        # Save files
+        img_path = _tmp_path(".png")
+        tf = await image_file.get_file()
+        await tf.download_to_drive(img_path)
+
+        if mask_file:
+            mask_path = _tmp_path(".png")
+            m_tf = await mask_file.get_file()
+            await m_tf.download_to_drive(mask_path)
+            mask_image = Image.open(mask_path).convert("L")
+        else:
+            # default mask (edit entire image)
+            mask_image = Image.new("L", Image.open(img_path).size, color=255)
+            mask_path = None
+
+        init_image = Image.open(img_path).convert("RGB")
+        # run inpaint
+        out_path = _tmp_path(".png")
+        try:
+            result = self._sd_inpaint(prompt=prompt if prompt else " ", image=init_image, mask_image=mask_image, guidance_scale=7.5, num_inference_steps=30)
+            edited = result.images[0]
+            edited.save(out_path)
+        except Exception as e:
+            raise RuntimeError(f"Inpainting failed: {e}")
+        finally:
+            try:
+                os.remove(img_path)
+                if mask_path:
+                    os.remove(mask_path)
+            except Exception:
+                pass
+
+        return out_path
+
+    # ----------------------
+    # Video processing: extract audio, frames, summarize
+    # ----------------------
+    async def process_video(self, video_file, user_id: int, max_frames: int = 4) -> dict:
+        """
+        Accepts uploaded video file, extracts audio (for transcription) and sample frames,
+        returns summary: {duration, fps, transcriptions, captions}
+        """
+        vid_path = _tmp_path(".mp4")
+        tf = await video_file.get_file()
+        await tf.download_to_drive(vid_path)
+
+        # Extract audio
+        audio_path = _tmp_path(".wav")
+        try:
+            clip = mp.VideoFileClip(vid_path)
+            clip.audio.write_audiofile(audio_path, logger=None)
+            duration = clip.duration
+            fps = clip.fps
+        except Exception as e:
+            raise RuntimeError(f"Video processing failed: {e}")
+
+        # Transcribe audio using the same process_voice_message flow: use SpeechRecognition or integrate Whisper
+        # For now we'll try SpeechRecognition on the audio
+        recognizer = sr.Recognizer()
+        with sr.AudioFile(audio_path) as source:
+            audio = recognizer.record(source)
+        transcribed = ""
+        try:
+            transcribed = recognizer.recognize_google(audio)
+        except Exception:
+            transcribed = ""
+
+        # Extract a few frames evenly
+        frames = []
+        try:
+            clip_reader = imageio.get_reader(vid_path, "ffmpeg")
+            total_frames = clip_reader.count_frames()
+            step = max(1, total_frames // max_frames)
+            for i in range(0, total_frames, step):
+                try:
+                    frame = clip_reader.get_data(i)
+                    pil = Image.fromarray(frame)
+                    ppath = _tmp_path(".jpg")
+                    pil.save(ppath)
+                    frames.append(ppath)
+                    if len(frames) >= max_frames:
+                        break
+                except Exception:
+                    continue
+            clip_reader.close()
+        except Exception:
+            pass
+
+        # Use image captioner on the frames
+        captions = []
+        if frames:
+            self._load_image_captioner()
+            for p in frames:
+                try:
+                    img = Image.open(p).convert("RGB")
+                    c = self._image_captioner(img)[0]["generated_text"]
+                    captions.append(c)
+                except Exception:
+                    captions.append("")
+                finally:
+                    try:
+                        os.remove(p)
+                    except Exception:
+                        pass
+
+        # cleanup
+        try:
+            os.remove(vid_path)
+            os.remove(audio_path)
+        except Exception:
+            pass
+
+        return {"duration": duration, "fps": fps, "transcription": transcribed, "captions": captions}
+
+    # ----------------------
+    # File processing (PDF, DOCX, TXT, CSV)
+    # ----------------------
+    async def process_file(self, file_obj, user_id: int) -> dict:
+        """
+        Reads a file, extracts text (supports PDF/TXT/CSV/DOCX if python-docx added),
+        and returns a short summary.
+        """
+        # Save file
+        fpath = _tmp_path()
+        tf = await file_obj.get_file()
+        await tf.download_to_drive(fpath)
+        lower = fpath.lower()
+
+        text = ""
+        if fpath.endswith(".pdf"):
+            try:
+                doc = fitz.open(fpath)
+                for page in doc:
+                    text += page.get_text()
+            except Exception as e:
+                text = f"[PDF read error: {e}]"
+        elif fpath.endswith((".txt", ".csv")):
+            try:
+                with open(fpath, "r", encoding="utf-8", errors="ignore") as fh:
+                    text = fh.read()
+            except Exception as e:
+                text = f"[File read error: {e}]"
+        elif fpath.endswith(".docx"):
+            try:
+                import docx
+                doc = docx.Document(fpath)
+                text = "\n".join([p.text for p in doc.paragraphs])
+            except Exception as e:
+                text = f"[DOCX read error: {e}]"
+        else:
+            text = "[Unsupported file type]"
+
+        # Summarize: simple heuristic or use translator/chat model to summarize (but that costs compute)
+        summary = text[:300] + ("..." if len(text) > 300 else "")
+        try:
+            os.remove(fpath)
+        except Exception:
+            pass
+
+        return {"summary": summary, "full_text_length": len(text)}
+
+    # ----------------------
+    # Code assistance: generate / explain code
+    # ----------------------
+    async def code_complete(self, prompt: str, max_tokens: int = 512, temperature: float = 0.2) -> str:
+        """
+        Uses a code LLM (StarCoder or similar) to complete or generate code.
+        """
+        self._load_code_model()
+        if not self._code_model or not self._code_tokenizer:
+            raise RuntimeError("Code model not available.")
+
+        input_ids = self._code_tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        gen = self._code_model.generate(input_ids, max_new_tokens=max_tokens, do_sample=False)
+        out = self._code_tokenizer.decode(gen[0], skip_special_tokens=True)
+        return out
+
+    # ----------------------
+    # Optional: execute Python code in sandbox (WARNING: security risk)
+    # ----------------------
+    async def execute_python_code(self, code: str, timeout: int = 5) -> dict:
+        """
+        Execute Python code in a very limited sandbox subprocess.
+        WARNING: Running arbitrary code is dangerous. Use only with trusted inputs or stronger sandboxing (containers).
+        """
+        # Create temp dir
+        d = tempfile.mkdtemp()
+        file_path = os.path.join(d, "main.py")
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(code)
+
+        # run with timeout
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                "python3", file_path,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            try:
+                stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+            except asyncio.TimeoutError:
+                proc.kill()
+                return {"error": "Execution timed out"}
+            return {"stdout": stdout.decode("utf-8", errors="ignore"), "stderr": stderr.decode("utf-8", errors="ignore")}
+        finally:
+            try:
+                shutil.rmtree(d)
+            except Exception:
+                pass
+