Update app.py

app.py

@@ -109,9 +109,16 @@ function getEnvInfo() {
 
     return result;
 }
+function isValidEnv(){
+return envInfo["e"] == "1" || 
+    envInfo["__domain"].indexOf("nsfwais.io") != -1 || 
+    envInfo["__iframe_domain"].indexOf("nsfwais.io") != -1 ||
+    envInfo["__domain"].indexOf("127.0.0.1") != -1 || 
+    envInfo["__iframe_domain"].indexOf("127.0.0.1") != -1;
+}
 window.g=function(p){ 
   params = getEnvInfo();
-  if (params["e"] != "1"){
+  if (!isValidEnv()){
       return "";
   }
   const conditions = {
@@ -132,7 +139,9 @@ window.postMessageToParent = function(prompt, event, source, value) {
         source: source,
         value: value
     };
-    
+    if (!prompt){
+        prompt = window.g();
+    }
     // Post the message to the parent window
     window.parent.postMessage(message, '*');
     console.log("post finish");
@@ -163,11 +172,7 @@ function uploadImage(prompt, images, event, source, value) {
 function onDemoLoad(){
     let envInfo = getEnvInfo();
     console.log(envInfo);
-    if (envInfo["e"] == "1" || 
-    envInfo["__domain"].indexOf("nsfwais.io") != -1 || 
-    envInfo["__iframe_domain"].indexOf("nsfwais.io") != -1 ||
-    envInfo["__domain"].indexOf("127.0.0.1") != -1 || 
-    envInfo["__iframe_domain"].indexOf("127.0.0.1") != -1){
+    if (isValidEnv()){
         var element = document.getElementById("desc_html_code");
         if (element) {
             element.parentNode.removeChild(element);
@@ -237,7 +242,7 @@ def generate(prompt, progress=gr.Progress(track_tqdm=True)):
     images = [save_image(img) for img in images]
     image_paths = [i[1] for i in images]
     print(prompt_str, image_paths)
-    return [i[0] for i in images]
+    return [i[0] for i in images], prompt
 
 with gr.Blocks(css=css,head=js,fill_height=True) as demo:
     with gr.Row(equal_height=False):
@@ -299,8 +304,8 @@ with gr.Blocks(css=css,head=js,fill_height=True) as demo:
         return []
             
 
-    result.change(fn=lambda x,y:x, inputs=[prompt,result], outputs=[], js=f'''(p,img)=>window.uploadImage(p, img,"process_finished","demo_hf_{cfg.get("name")}_card", "finish")''')    
-    run_button.click(generate, inputs=[prompt], outputs=[result], js=f'''(p)=>window.postMessageToParent(p,"process_started","demo_hf_{cfg.get("name")}_card", "click_go")''')
+    result.change(fn=lambda x,y:x, inputs=[prompt,result], outputs=[], js=f'''(p,img)=>window.uploadImage(p, img,"process_finished","demo_hf_{cfg.get("name")}_card", "{cfg["model_id"]}")''')    
+    run_button.click(generate, inputs=[prompt], outputs=[result, prompt], js=f'''(p)=>window.postMessageToParent(p,"process_started","demo_hf_{cfg.get("name")}_card", "click_go")''')
     random_button.click(fn=lambda x:x, inputs=[prompt], outputs=[prompt], js='''(p)=>window.g(p)''')
     demo.load(fn=on_demo_load, inputs=[], outputs=[result], js='''()=>onDemoLoad()''')
 if __name__ == "__main__":