Update app.py

app.py

@@ -1,54 +1,44 @@
 import gradio as gr  
 from transformers import AutoModel, AutoTokenizer  
-import torch
-import torchvision.transforms as T
-from torchvision.transforms.functional import InterpolationMode
-from Models.modeling_llavaqw import LlavaQwModel
-
-
-
-
-IMAGENET_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_STD = (0.229, 0.224, 0.225)
+import torch  
+import torchvision.transforms as T  
+from torchvision.transforms.functional import InterpolationMode  
+from Models.modeling_llavaqw import LlavaQwModel  
 
+IMAGENET_MEAN = (0.485, 0.456, 0.406)  
+IMAGENET_STD = (0.229, 0.224, 0.225)  
 
 model_name = "torettomarui/Llava-qw"  
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)
-model = LlavaQwModel.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-).to(torch.bfloat16).eval()#.cuda()
-
-def build_transform(input_size):
-    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
-    transform = T.Compose([
-        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
-        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
-        T.ToTensor(),
-        T.Normalize(mean=MEAN, std=STD)
-    ])
-    return transform
-
-
-def preprocess_image(file_path, image_size=448):
-
-    transform = build_transform(image_size)
-    pixel_values = transform(file_path)
-    return torch.stack([pixel_values]).to(torch.bfloat16)#.cuda()
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)  
+model = LlavaQwModel.from_pretrained(  
+    model_name,  
+    torch_dtype=torch.bfloat16,  
+    trust_remote_code=True,  
+).to(torch.bfloat16).eval().cuda()  
+
+def build_transform(input_size):  
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD  
+    transform = T.Compose([  
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),  
+        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),  
+        T.ToTensor(),  
+        T.Normalize(mean=MEAN, std=STD)  
+    ])  
+    return transform  
+
+def preprocess_image(file_path, image_size=448):  
+    transform = build_transform(image_size)  
+    pixel_values = transform(file_path)  
+    return torch.stack([pixel_values]).to(torch.bfloat16).cuda()  
 
 def generate_response(image, text):  
-
-    pixel_values = preprocess_image(image)
-
-    generation_config = dict(max_new_tokens=2048, do_sample=False)
-
-    question = '<image>\n' + text
-
-    response = model.chat(tokenizer, pixel_values, question, generation_config)
-    
+    pixel_values = preprocess_image(image)  
+    generation_config = dict(max_new_tokens=2048, do_sample=False)  
+    question = '<image>\n' + text  
+    response = model.chat(tokenizer, pixel_values, question, generation_config)  
     return response  
 
+# 添加示例图像和文本  
 examples = [  
     ["./text.png", "图中的文字是什么?"],  
 ]