Update app.py

app.py

@@ -1,36 +1,55 @@
 import gradio as gr
 
-def generate_text(title, context):
-    return f"Title:{title}\nContext:{context}\n..."
+def generate_multi_model(input_img):
+    output_img = input_img
+    return output_img
 
-def generate_mutimodal(title, context, img):
+def train_one_shot(title, context, img):
+    return f"Title:{title}\nContext:{context}\n...{img}"
+
+def generate_one_shot(title, context, img):
+    return f"Title:{title}\nContext:{context}\n...{img}"
+
+def train_zero_shot(title, context, img):
+    return f"Title:{title}\nContext:{context}\n...{img}"
+
+def generate_zero_shot(title, context, img):
     return f"Title:{title}\nContext:{context}\n...{img}"
 
 with gr.Blocks() as demo:
     # 顶部文字
-    gr.Markdown("# 多模态可控XX生成")
+    gr.Markdown("# MMFS")
 
     # 多个tab
     with gr.Tabs():
 
-        with gr.TabItem("新闻输入"):
-            title = gr.Textbox(label="标题", lines=1, placeholder="请输入标题")
-            context = gr.Textbox(label="正文", lines=2, placeholder="请输入正文")
-            text_button = gr.Button("生成")
+        with gr.TabItem("multi-model"):
+            multi_input_img = gr.Image(shape=(200, 200),  label="请上传人像图片")
+            multi_model_button = gr.Button("随机风格化")
             
-            text_output = gr.Textbox(label="生成内容", lines=5, placeholder="生成列表")
+            multi_output_img = gr.Image(shape=(200, 200),  label="人像风格化图")
+
+
+        with gr.TabItem("one-shot"):
+            one_shot_src_img = gr.Image(shape=(200, 200),  label="请上传内容图片")
+            one_shot_ref_img = gr.Image(shape=(200, 200),  label="请上传风格图片")
+            one_shot_train_button = gr.Button("训练模型")
+            one_shot_test_button = gr.Button("风格化")
 
+            one_shot_output_img = gr.Image(shape=(200, 200),  label="one-shot风格化图")
 
-        with gr.TabItem("图文输入"):
-            title = gr.Textbox(label="标题", lines=1, placeholder="请输入标题")
-            context = gr.Textbox(label="正文", lines=2, placeholder="请输入正文")
-            img = gr.Image(shape=(200, 200),  label="请上传图片")
-            img_button = gr.Button("生成")
+        
+        with gr.TabItem("zero-shot"):
+            zero_shot_src_img = gr.Image(shape=(200, 200),  label="请上传内容图片")
+            zero_shot_ref_prompt = gr.Textbox(label="prompt", lines=1, placeholder="请输入参考风格描述（英文）")
+            zero_shot_train_button = gr.Button("训练模型")
+            zero_shot_test_button = gr.Button("风格化")
 
-            img_output = gr.Textbox(label="生成内容", lines=5, placeholder="生成列表")
+            zero_shot_output_img = gr.Image(shape=(200, 200),  label="zero-shot风格化图")
 
 
-    text_button.click(fn=generate_text, inputs=[title,context], outputs=text_output)
-    img_button.click(fn=generate_mutimodal, inputs=[title,context,img], outputs=img_output)
+    multi_model_button.click(fn=generate_multi_model, inputs=multi_input_img, outputs=multi_output_img)
+    one_shot_test_button.click(fn=generate_one_shot, inputs=[title,context,img], outputs=img_output)
+    zero_shot_test_button.click(fn=generate_zero_shot, inputs=[title,context,img], outputs=img_output)
 
 demo.launch()