Full ARK-AI multi-modal setup

app.py

@@ -1,126 +1,79 @@
-# app.py
-import gradio as gr
-from transformers import pipeline
-
-# =============================
-# LOAD YOUR MODELS
-# =============================
-
-# ----- Text models -----
-text_models = [
-    pipeline("text-generation", model="openai/gpt-oss-120b"),
-    pipeline("text-generation", model="deepseek-ai/DeepSeek-V3.1-Base"),
-    pipeline("text-generation", model="zai-org/GLM-4.5"),
-    pipeline("text-generation", model="zai-org/chatglm3-6b"),
-    pipeline("text-generation", model="openai/gpt-oss-20b")
-]
-
-# ----- Audio models -----
-tts_models = [
-    pipeline("text-to-speech", model="sk0032/coqui-tts-model"),
-    pipeline("text-to-speech", model="sysf/Edge-TTS")
-]
-
-# ----- Whisper ASR -----
-speech_to_text_model = pipeline("automatic-speech-recognition", model="openai/whisper-large-v3-turbo")
-
-# ----- Image models -----
-image_models = [
-    pipeline("text-to-image", model="GD-ML/USP-Image_Generation"),
-    pipeline("text-to-image", model="Qwen/Qwen-Image"),
-    pipeline("text-to-image", model="Comfy-Org/Qwen-Image-DiffSynth-ControlNets"),
-    pipeline("image-to-text", model="Salesforce/blip2-opt-2.7b"),
-    pipeline("vision-to-text", model="Vision-CAIR/MiniGPT-4")
-]
-
-# ----- Video models -----
-video_models = [
-    pipeline("text-to-video", model="sudip1987/Generate_videos_with_Veo3"),
-    pipeline("text-to-video", model="ali-vilab/text-to-video-ms-1.7b")
-]
-
-# ----- Music model -----
-music_model = pipeline("text-to-music", model="sander-wood/text-to-music")
-
-# =============================
-# HELPER FUNCTIONS
-# =============================
-
-def merge_text_models(prompt):
-    responses = []
-    for model in text_models:
-        try:
-            out = model(prompt)[0]['generated_text']
-            responses.append(out)
-        except Exception as e:
-            responses.append(f"[Model failed: {str(e)}]")
-    return "\n---\n".join(responses)
-
-def generate_audio(text):
-    audio_files = []
-    for model in tts_models:
-        try:
-            audio_path = model(text)
-            audio_files.append(audio_path)
-        except Exception as e:
-            audio_files.append(None)
-    return audio_files[0] if audio_files else None
-
-def generate_images(text):
-    imgs = []
-    for model in image_models:
-        try:
-            img = model(text)[0]['image']
-            imgs.append(img)
-        except Exception as e:
-            continue
-    return imgs[:3]  # Show top 3 images
-
-def generate_videos(text):
-    vids = []
-    for model in video_models:
-        try:
-            vid = model(text)
-            vids.append(vid)
-        except Exception as e:
-            continue
-    return vids[:1]  # Show one video
-
-# =============================
-# MAIN ARK-AI FUNCTION
-# =============================
-
-def ark_ai_main(prompt):
-    # Text
-    text_output = merge_text_models(prompt)
-    
-    # Inject personality
-    personality = "ARK-AI (fun, savage, chaotic-good) says:\n"
-    full_text = personality + text_output
+# Backend.py
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+app = Flask(__name__)
+CORS(app)
+
+# List of models
+MODEL_PATHS = {
+    "gpt_oss_120b": "openai/gpt-oss-120b",
+    "deepseek_v3": "deepseek-ai/DeepSeek-V3.1-Base",
+    "gemini_25_pro": "afu4642tD/gemini-2.5-pro",
+    "veo3": "sudip1987/Generate_videos_with_Veo3",
+    "open_sora": "hpcai-tech/Open-Sora-v2",
+    "usp_image": "GD-ML/USP-Image_Generation",
+    "text_to_music": "sander-wood/text-to-music",
+    "qwen_image": "Qwen/Qwen-Image",
+    "qwen_image_diff": "Comfy-Org/Qwen-Image-DiffSynth-ControlNets",
+    "coqui_tts": "sk0032/coqui-tts-model",
+    "edge_tts": "sysf/Edge-TTS",
+    "whisper_large": "openai/whisper-large-v3-turbo",
+    "blip2_opt": "Salesforce/blip2-opt-2.7b",
+    "mini_gpt4": "Vision-CAIR/MiniGPT-4",
+    "glm_45": "zai-org/GLM-4.5",
+    "chatglm3": "zai-org/chatglm3-6b",
+    "gpt_oss_20b": "openai/gpt-oss-20b",
+    "m2m100": "facebook/m2m100_1.2B",
+    "tiny_marian": "onnx-community/tiny-random-MarianMTModel",
+    "memory_transformer": "Grpp/memory-transformer-ru",
+    "rl_memory_agent": "BytedTsinghua-SIA/RL-MemoryAgent-14B",
+    "m3_agent": "ByteDance-Seed/M3-Agent-Memorization",
+    "text_to_video": "ali-vilab/text-to-video-ms-1.7b"
+}
+
+def generate_single_answer(prompt, model_key):
+    """Load model, generate answer, free memory"""
+    model_name = MODEL_PATHS[model_key]
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        device_map="auto" if torch.cuda.is_available() else None,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        output = model.generate(**inputs, max_new_tokens=200)
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Clean up to save memory
+    del model, tokenizer, inputs, output
+    torch.cuda.empty_cache()
+
+    return answer
+
+@app.route("/ask", methods=["POST"])
+def ask():
+    data = request.json
+    prompt = data.get("prompt", "")
+    selected_models = data.get("models", ["gpt_oss_120b", "deepseek_v3", "gemini_25_pro"])
     
-    # Media
-    image_output = generate_images(prompt)
-    video_output = generate_videos(prompt)
-    audio_output = generate_audio(prompt)
-    
-    return full_text, image_output, video_output, audio_output
-
-# =============================
-# GRADIO INTERFACE
-# =============================
-
-iface = gr.Interface(
-    fn=ark_ai_main,
-    inputs=gr.Textbox(lines=3, placeholder="Ask ARK-AI anything..."),
-    outputs=[
-        gr.Textbox(label="ARK-AI Text Response"),
-        gr.Gallery(label="Images Generated"),
-        gr.Video(label="Video Generated"),
-        gr.Audio(label="Audio Response")
-    ],
-    title="ARK-AI Multi-Modal Assistant",
-    description="ARK-AI: Savage, funny, chaotic-good AI assistant merging text, image, audio, and video models.",
-    css="styles.css"  # Optional: liquid-glass UI
-)
-
-iface.launch()
+    # Merge answers from all selected models
+    answers = []
+    for model_key in selected_models:
+        if model_key in MODEL_PATHS:
+            try:
+                ans = generate_single_answer(prompt, model_key)
+                answers.append(ans)
+            except Exception as e:
+                answers.append(f"[Error loading {model_key}]")
+
+    # Return one merged answer
+    final_answer = " | ".join(answers)
+    return jsonify({"answer": final_answer})
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000, debug=True)