Update app.py

app.py

@@ -79,7 +79,38 @@ footer {display: none !important;}
 js = '''
 <script src="https://huggingface.co/spaces/nsfwalex/sd_card/resolve/main/prompt.js"></script>
 <script>
+function getEnvInfo() {
+    const result = {};
+
+    // Get URL parameters
+    const urlParams = new URLSearchParams(window.location.search);
+    for (const [key, value] of urlParams) {
+        result[key] = value;
+    }
+
+    // Get current domain and convert to lowercase
+    result["__domain"] = window.location.hostname.toLowerCase();
+
+    // Get iframe parent domain, if any, and convert to lowercase
+    try {
+        if (window.self !== window.top) {
+            result["__iframe_domain"] = document.referrer 
+                ? new URL(document.referrer).hostname.toLowerCase()
+                : "unable to get iframe parent domain";
+        }else{
+            result["__iframe_domain"] = null;
+        }
+    } catch (e) {
+        result["__iframe_domain"] = "unable to access iframe parent domain";
+    }
+
+    return result;
+}
 window.g=function(){ 
+  params = getUrlParamsAndDomain()
+  if (params["e"] != "1"){
+      return "";
+  }
   const conditions = {
     "tag": ["normal", "sexy", "porn"],
     "exclude_category": ["Clothing"],
@@ -89,20 +120,7 @@ window.g=function(){
   console.log(prompt);
   return prompt
 }
-function checkDomain(img, str) {
-  // Get the current page's hostname (domain)
-  const currentDomain = window.location.hostname;
-  
-  // Convert both the domain and the input string to lowercase for case-insensitive comparison
-  const lowerDomain = currentDomain.toLowerCase();
-  const lowerStr = str.toLowerCase();
-  
-  // Check if the domain contains the string
-  if lowerDomain.includes(lowerStr){
-      return null;
-  }
-  return img;
-}
+
 window.postMessageToParent = function(prompt, event, source, value) {
     // Construct the message object with the provided parameters
     console.log("post start",event, source, value);
@@ -138,6 +156,10 @@ function uploadImage(prompt, images, event, source, value) {
         console.log("Not in an iframe, can't post to parent");
     }
 }
+function onDemoLoad(){
+    let envInfo = getEnvInfo()
+    return envInfo["__domain"], envInfo["__iframe_domain"]
+}
 </script>
 '''
 def save_image(img):
@@ -220,20 +242,19 @@ with gr.Blocks(css=css,head=js,fill_height=True) as demo:
                 random_button = gr.Button("Surprise Me", scale=1, min_width=10)
             run_button = gr.Button( "GO!", scale=1, min_width=20, variant="primary",icon="https://huggingface.co/spaces/nsfwalex/sd_card/resolve/main/hot.svg")
             
-    def on_demo_load(request: gr.Request):
+    def on_demo_load(request: gr.Request, domain, iframe_domain):
         params = dict(request.query_params)
-        print("load_demo, url params", params, "host", request.client.host)
+        print("load_demo, url params", params, "domain", domain, "iframe", iframe_domain)
         if params.get("e", "0") == "1":
             #update the image
             #bind events
-            random_button.click(fn=lambda x:x, inputs=[prompt], outputs=[prompt], js='''()=>window.g()''')
             return Image.open(default_image)
         return None
             
 
     result.change(fn=lambda x: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")''')
-            
-    demo.load(fn=on_demo_load, inputs=[], outputs=[result])#, js='''(img)=>checkDomain(img, "huggingface")''')
+    random_button.click(fn=lambda x:x, inputs=[prompt], outputs=[prompt], js='''()=>window.g()''')
+    demo.load(fn=on_demo_load, inputs=[], outputs=[result], js='''()=>onDemoLoad()''')
 if __name__ == "__main__":
     demo.queue().launch(show_api=False)