Update some files

app.py

@@ -1,68 +1,67 @@
-import os
-import gradio as gr
-import spaces
+import math
 import torch
-import tempfile
-import imageio
-from decord import VideoReader, cpu
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+import gradio as gr
 from PIL import Image
 
-hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-model_id = "google/gemma-3-27b-it"
-NUM_FRAMES = 8
-
-# 从视频中采样 N 帧
-def sample_video_frames(video_path, num_frames=8):
-    vr = VideoReader(video_path, ctx=cpu(0))
-    total_frames = len(vr)
-    indices = [int(i) for i in torch.linspace(0, total_frames - 1, steps=num_frames)]
-    
-    # 关键点：强制转换为 PIL.Image
-    frames = [Image.fromarray(vr[i].asnumpy()) for i in indices]
-    return frames
-
-# 推理函数：加载模型、采样视频帧、推理
-@spaces.GPU
-def analyze_video(video_file):
-    # video_file 是路径字符串
-    frames = sample_video_frames(video_file)
-
-    # 构造 prompt
-    system_prompt = (
-        "You are a helpful AI assistant that analyzes AR effects in videos. "
-        "Evaluate the realism and placement of virtual objects in the provided video frames."
-    )
-    user_prompt = "Based on the frames, describe how well the AR objects blend into the real environment."
+# === 分配层到多 GPU ===
+def split_model(model_path):
+    from transformers import AutoConfig
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    num_layers = config.llm_config.num_hidden_layers
+    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
+    num_layers_per_gpu = [num_layers_per_gpu] * world_size
+    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
+    layer_cnt = 0
+    for i, num_layer in enumerate(num_layers_per_gpu):
+        for _ in range(num_layer):
+            device_map[f'language_model.model.layers.{layer_cnt}'] = i
+            layer_cnt += 1
+    device_map['vision_model'] = 0
+    device_map['mlp1'] = 0
+    device_map['language_model.model.tok_embeddings'] = 0
+    device_map['language_model.model.embed_tokens'] = 0
+    device_map['language_model.output'] = 0
+    device_map['language_model.model.norm'] = 0
+    device_map['language_model.model.rotary_emb'] = 0
+    device_map['language_model.lm_head'] = 0
+    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
+    return device_map
 
-    history = [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": system_prompt}]
-        },
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": user_prompt}] + [{"type": "image", "image": frame} for frame in frames]
-        }
-    ]
+# === 模型路径 ===
+model_path = "OpenGVLab/InternVL3-14B"
+device_map = split_model(model_path)
 
-    pipe = pipeline(
-        "image-text-to-text",
-        model=model_id,
-        token=hf_token,
-        torch_dtype=torch.bfloat16,
-        model_kwargs={"device_map": "auto"}
-    )
+# === 加载模型和处理器 ===
+model = AutoModel.from_pretrained(
+    model_path,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True,
+    device_map=device_map
+).eval()
 
-    result = pipe(text=history, max_new_tokens=512)
-    return result[0]["generated_text"][-1]["content"]
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
 
+# === 推理函数 ===
+def infer(image: Image.Image, prompt: str):
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_new_tokens=512)
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    return answer
 
-# Gradio 界面
+# === Gradio 界面 ===
 gr.Interface(
-    fn=analyze_video,
-    inputs=gr.Video(label="Upload an AR Video (.mp4 only)"),
-    outputs=gr.Textbox(label="Gemma Analysis Result"),
-    title="Gemma-3-27B Video Analysis (ZeroGPU)",
-    description="Uploads a video, extracts 8 frames, and uses Gemma-3-27B to analyze AR realism."
+    fn=infer,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Your Prompt", placeholder="Ask a question about the image...")
+    ],
+    outputs="text",
+    title="InternVL3-14B Multimodal Demo",
+    description="Upload an image and ask a question. InternVL3-14B will answer using vision + language."
 ).launch()