Update app.py

Add button for different parts of training model process

app.py

@@ -1,3 +1,16 @@
+import argparse
+import itertools
+import math
+import os
+from pathlib import Path
+from typing import Optional
+import subprocess
+import sys
+
+import torch
+
+from spanish_medica_llm import run_training
+
 import gradio as gr
 
 #def greet(name):
@@ -6,15 +19,76 @@ import gradio as gr
 #iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 #iface.launch()
 
-def update(name):
+def generate_model(name):
     return f"Welcome to Gradio, {name}!"
+    
+def generate(prompt):
+    from diffusers import StableDiffusionPipeline
+    
+    pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
+    pipe = pipe.to("cuda")
+    image = pipe(prompt).images[0]  
+    return(image)
+    
+def evaluate_model():
+    #from diffusers import StableDiffusionPipeline
+    
+    #pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
+    #pipe = pipe.to("cuda")
+    #image = pipe(prompt).images[0]  
+    return("Evaluate Model")
+
+def train_model(*inputs):
+    if "IS_SHARED_UI" in os.environ:
+        raise gr.Error("This Space only works in duplicated instances")
+  
+    args_general = argparse.Namespace(
+                image_captions_filename = True,
+                train_text_encoder = True,
+                stop_text_encoder_training = stptxt,
+                save_n_steps = 0,
+                pretrained_model_name_or_path = model_to_load,
+                instance_data_dir="instance_images",
+                class_data_dir=class_data_dir,
+                output_dir="output_model",
+                instance_prompt="",
+                seed=42,
+                resolution=512,
+                mixed_precision="fp16",
+                train_batch_size=1,
+                gradient_accumulation_steps=1,
+                use_8bit_adam=True,
+                learning_rate=2e-6,
+                lr_scheduler="polynomial",
+                lr_warmup_steps = 0,
+                max_train_steps=Training_Steps,     
+    )
+    run_training(args_general)
+    torch.cuda.empty_cache()
+    #convert("output_model", "model.ckpt")
+    #shutil.rmtree('instance_images')
+    #shutil.make_archive("diffusers_model", 'zip', "output_model")
+    #with zipfile.ZipFile('diffusers_model.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
+    #    zipdir('output_model/', zipf)
+    torch.cuda.empty_cache()
+    return [gr.update(visible=True, value=["diffusers_model.zip"]), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)]
+
+def stop_model(*input):
+    return f"Model with Gradio!"
+
 
 with gr.Blocks() as demo:
     gr.Markdown("Start typing below and then click **Run** to see the output.")
     with gr.Row():
         inp = gr.Textbox(placeholder="What is your name?")
         out = gr.Textbox()
-    btn = gr.Button("Run")
-    btn.click(fn=update, inputs=inp, outputs=out)
+    btn_response = gr.Button("Generate Response")
+    btn_response.click(fn=generate_model, inputs=inp, outputs=out)
+    btn_train = gr.Button("Train Model")
+    btn_train.click(fn=train_model, inputs=[], outputs=out)
+    btn_evaluate = gr.Button("Evaluate Model")
+    btn_evaluate.click(fn=evaluate_model, inputs=[], outputs=out)
+    btn_stop = gr.Button("Stop Model")
+    btn_stop.click(fn=stop_model, inputs=[], outputs=out)
 
 demo.launch()