Create app.py

app.py

@@ -0,0 +1,107 @@
+import gradio as gr
+from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+vision_pipe = pipeline(
+    "image-to-text",
+    model="nlpconnect/vit-gpt2-image-captioning",
+    device=0 if device == "cuda" else -1
+)
+
+text_pipe = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-base",
+    max_length=256,
+    device=0 if device == "cuda" else -1
+)
+
+sd_model_id = "runwayml/stable-diffusion-v1-5"
+sd_pipe = StableDiffusionPipeline.from_pretrained(
+    sd_model_id,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+sd_pipe = sd_pipe.to(device)
+
+def analyze_and_enhance(image, style, language):
+    if image is None:
+        return None, "请先上传图表。", "", None, "", ""
+
+    raw_caption = vision_pipe(image)[0]["generated_text"]
+
+    prompt_academic = (
+        "Rewrite the following description in a formal academic style, "
+        "suitable for a scientific paper. Be concise but precise.\n\n"
+        f"Description: {raw_caption}"
+    )
+    academic_desc = text_pipe(prompt_academic)[0]["generated_text"]
+
+    caption_prompt = (
+        "Write a 1-2 sentence figure caption for an academic paper, "
+        "based on this description:\n\n"
+        f"{academic_desc}\n\nCaption:"
+    )
+    caption_text = text_pipe(caption_prompt)[0]["generated_text"]
+
+    summary_prompt = (
+        "Write a short 3-4 sentence paragraph that explains the key trend "
+        "and message of this figure for the results section of a paper.\n\n"
+        f"{academic_desc}\n\nParagraph:"
+    )
+    summary_text = text_pipe(summary_prompt)[0]["generated_text"]
+
+    if language == "中文":
+        caption_text = text_pipe(
+            f"Translate this figure caption into Chinese:\n{caption_text}"
+        )[0]["generated_text"]
+        summary_text = text_pipe(
+            f"Translate this paragraph into Chinese:\n{summary_text}"
+        )[0]["generated_text"]
+        academic_desc_out = text_pipe(
+            f"Translate this academic description into Chinese:\n{academic_desc}"
+        )[0]["generated_text"]
+    else:
+        academic_desc_out = academic_desc
+
+    sd_prompt = (
+        f"{style} infographic style illustration of: {academic_desc} "
+        "minimal, clean, flat colors, suitable for a scientific slide."
+    )
+
+    with torch.autocast(device if device == "cuda" else "cpu"):
+        enhanced_image = sd_pipe(sd_prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
+
+    return image, raw_caption, academic_desc_out, enhanced_image, caption_text, summary_text
+
+with gr.Blocks() as demo:
+    gr.Markdown("# ChartSmith – AI 论文图表生成助手")
+
+    with gr.Row():
+        with gr.Column():
+            img_in = gr.Image(type="pil", label="上传你的学术图表（截图也可以）")
+            style = gr.Dropdown(
+                ["Formal academic", "Infographic", "Magazine-style"],
+                value="Formal academic",
+                label="重绘风格"
+            )
+            language = gr.Radio(["English", "中文"], value="English", label="输出语言")
+            btn = gr.Button("分析并美化图表")
+        with gr.Column():
+            orig_img = gr.Image(label="原始图表")
+            raw_caption_box = gr.Textbox(label="Step 2: 初步自动描述（Vision-LLM）")
+            academic_box = gr.Textbox(label="Step 3: 学术化解释（Academic Explanation）")
+            enhanced_img = gr.Image(label="Step 4: 美化 / 重绘后的图示")
+            caption_box = gr.Textbox(label="Step 5: 自动生成图注（Caption）")
+            summary_box = gr.Textbox(label="Step 5: 图表相关简短摘要 / 讨论")
+
+    btn.click(
+        analyze_and_enhance,
+        inputs=[img_in, style, language],
+        outputs=[orig_img, raw_caption_box, academic_box, enhanced_img, caption_box, summary_box]
+    )
+
+if __name__ == "__main__":
+    demo.launch()