Update app.py

app.py

@@ -1,103 +1,82 @@
-# import gradio as gr
-# import spaces
-# from PIL import Image
-# import torch
-# from transformers import AutoModelForCausalLM, AutoProcessor
-
-
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True).to(device)
-# processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True)
-
-# @spaces.GPU(duration=180)
-# def infer(prompt, image):
-#     max_size = 256
-#     width, height = image.size
-#     if width > height:
-#         new_width = max_size
-#         new_height = int((new_width / width) * height)
-#     else:
-#         new_height = max_size
-#         new_width = int((new_height / height) * width)
-
-#     image = image.resize((new_width, new_height), Image.LANCZOS)
-
-#     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
-
-#     generated_ids = model.generate(
-#         input_ids=inputs["input_ids"],
-#         pixel_values=inputs["pixel_values"],
-#         max_new_tokens=1024,
-#         do_sample=False,
-#         num_beams=3
-#     )
+import gradio as gr
+import spaces
+from PIL import Image
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True).to(device)
+processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True)
+
+@spaces.GPU(duration=180)
+def infer(prompt, image):
+    max_size = 256
+    width, height = image.size
+    if width > height:
+        new_width = max_size
+        new_height = int((new_width / width) * height)
+    else:
+        new_height = max_size
+        new_width = int((new_height / height) * width)
+
+    image = image.resize((new_width, new_height), Image.LANCZOS)
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        do_sample=False,
+        num_beams=3
+    )
     
-#     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
 
-#     parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
+    parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
     
-#     return parsed_answer
-
-# css = """
-# #col-container {
-#     margin: 0 auto;
-#     max-width: 800px;
-# }
-# """
-
-# with gr.Blocks(css=css) as app:
-#     with gr.Column(elem_id="col-container"):
-#         gr.Markdown(f"""# Tag The Image
-#         Get tag based on images using the Florence-2-base-PromptGen-v1.5 model.
-#         """)
-
-#         with gr.Row():
-#             prompt = gr.Text(
-#                 label="Prompt",
-#                 show_label=False,
-#                 max_lines=1,
-#                 placeholder="Enter your prompt or blank here.",
-#                 container=False,
-#             )
-#             image_input = gr.Image(
-#                 label="Image",
-#                 type="pil",
-#                 show_label=False,
-#                 container=False,
-#             )
-#             run_button = gr.Button("Run", scale=0)
-
-#         result = gr.Textbox(label="Generated Text", show_label=False)
-
-
-#     gr.on(
-#         triggers=[run_button.click, prompt.submit],
-#         fn=infer,
-#         inputs=[prompt, image_input],
-#         outputs=[result]
-#     )
-
-# app.queue()
-# app.launch(show_error=True)
-
-from fastapi import FastAPI
-import gradio as gr
-
-CUSTOM_PATH = "/gradio"
-
-app = FastAPI()
-
-@app.get("/")
-def read_main():
-    return {"message": "This is your main app"}
-
-# 新增的 /test 路由
-@app.get("/test")
-def test_api():
-    return {"message": "This is the test API"}
-
-io = gr.Interface(lambda x: "Hello, " + x + "!", "textbox", "textbox")
-app = gr.mount_gradio_app(app, io, path=CUSTOM_PATH)
-
-# 运行这个 FastAPI 应用：`uvicorn run:app`，并访问 http://localhost:8000/gradio 以及 http://localhost:8000/test
+    return parsed_answer
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 800px;
+}
+"""
+
+with gr.Blocks(css=css) as app:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# Tag The Image
+        Get tag based on images using the Florence-2-base-PromptGen-v1.5 model.
+        """)
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt or blank here.",
+                container=False,
+            )
+            image_input = gr.Image(
+                label="Image",
+                type="pil",
+                show_label=False,
+                container=False,
+            )
+            run_button = gr.Button("Run", scale=0)
+
+        result = gr.Textbox(label="Generated Text", show_label=False)
+
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, image_input],
+        outputs=[result]
+    )
+
+app.queue()
+app.launch(show_error=True)