Update app.py

app.py

@@ -24,29 +24,25 @@ def init_db(file='logs.db'):
     conn.close()
 
 # Log a request
-def log_request(prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo, huggingface_api_key, use_unhinged):
+def log_request(prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo, huggingface_api_key):
     log_message = f"Request: prompt='{prompt}', is_negative={is_negative}, steps={steps}, cfg_scale={cfg_scale}, "
     log_message += f"sampler='{sampler}', seed={seed}, strength={strength}, use_dev={use_dev}, "
     log_message += f"enhance_prompt_style='{enhance_prompt_style}', enhance_prompt_option={enhance_prompt_option}, "
-    log_message += f"nemo_enhance_prompt_style='{nemo_enhance_prompt_style}', use_mistral_nemo={use_mistral_nemo}, "
-    log_message += f"use_unhinged={use_unhinged}"
+    log_message += f"nemo_enhance_prompt_style='{nemo_enhance_prompt_style}', use_mistral_nemo={use_mistral_nemo}"
     if huggingface_api_key:
-        log_message += f", huggingface_api_key='{huggingface_api_key}'"
+        log_message += f"huggingface_api_key='{huggingface_api_key}'"
     
     conn = sqlite3.connect('acces_log.log')
     c = conn.cursor()
     c.execute("INSERT INTO logs VALUES (?, ?)", (datetime.now().isoformat(), log_message))
     conn.commit()
     conn.close()
-
+    
 # os.makedirs('assets', exist_ok=True)
 if not os.path.exists('icon.png'):
     os.system("wget -O icon.png https://i.pinimg.com/564x/64/49/88/644988c59447eb00286834c2e70fdd6b.jpg")
-    
 API_URL_DEV = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
-API_URL_UNHINGED = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"
-
 timeout = 100
 init_db('acces_log.log')
 
@@ -55,14 +51,6 @@ logging.basicConfig(filename='access.log', level=logging.INFO,
                     format='%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
 
 
-# Function to query the uncensored API
-def query_unhinged_api(prompt):
-    key = os.getenv('HUGGINGFACEHUB_API_TOKEN')
-    headers = {"Authorization": f"Bearer {key}"}
-    payload = {"inputs": prompt}
-    response = requests.post(API_URL_UNHINGED, headers=headers, json=payload)
-    return response.content
-
 
 def log_requestold(prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo, huggingface_api_key):
     log_message = f"Request: prompt='{prompt}', is_negative={is_negative}, steps={steps}, cfg_scale={cfg_scale}, "
@@ -153,17 +141,11 @@ def mistral_nemo_call(prompt, API_TOKEN, model="mistralai/Mistral-Nemo-Instruct-
         response += message.choices[0].delta.content
     return response
     
-def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key=None, use_dev=False, enhance_prompt_style="generic", enhance_prompt_option=False, nemo_enhance_prompt_style="generic", use_mistral_nemo=False, use_unhinged=False):
-    
-    log_request(prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo, huggingface_api_key, use_unhinged)
+def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key=None, use_dev=False,enhance_prompt_style="generic", enhance_prompt_option=False, nemo_enhance_prompt_style="generic", use_mistral_nemo=False):
     
+    log_request(prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo, huggingface_api_key)
     # Determine which API URL to use
-    if use_unhinged:
-        api_url = API_URL_UNHINGED
-    elif use_dev:
-        api_url = API_URL_DEV
-    else:
-        api_url = API_URL
+    api_url = API_URL_DEV if use_dev else API_URL
 
     # Check if the request is an API call by checking for the presence of the huggingface_api_key
     is_api_call = huggingface_api_key is not None
@@ -184,6 +166,7 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
 
     key = random.randint(0, 999)
     prompt = check_ubuse(prompt)
+    #prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
     print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
 
     original_prompt = prompt
@@ -201,17 +184,6 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     if seed == -1:
         seed = random.randint(1, 1000000000)
 
-    # Use uncensored API if selected
-    if use_unhinged:
-        unhinged_gallery = 'unhinged_images'
-        os.makedirs(unhinged_gallery, exist_ok=True)
-        image_bytes = query_unhinged_api(prompt)
-        image = Image.open(io.BytesIO(image_bytes))
-        output_path = f"./{unhinged_gallery}/unhinged_output_{random.randint(0, 999)}.png"
-        print(output_path)
-        image.save(output_path)
-        return output_path, seed, final_prompt
-    
     payload = {
         "inputs": final_prompt,
         "is_negative": is_negative,
@@ -242,7 +214,6 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     except Exception as e:
         print(f"Error when trying to open the image: {e}")
         return None, None, None
-        
 
 css = """
 body {
@@ -305,6 +276,7 @@ css1 = """
     font-weight: bold;
 }
 """
+
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
     gr.HTML("""
         <center>
@@ -333,8 +305,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                         enhance_prompt_option = gr.Checkbox(label="Enhance Prompt", value=False, elem_id="enhance-prompt-checkbox")
                         use_mistral_nemo = gr.Checkbox(label="Use Mistral Nemo", value=False, elem_id="use-mistral-checkbox")
                         nemo_prompt_style =  gr.Textbox(label="Nemo Enhance Prompt Style", placeholder="Enter style for the prompt enhancer here", elem_id="nemo-enhance-prompt-style")
-                        use_unhinged = gr.Checkbox(label="Use Flux Unhinged API", value=False, elem_id="use-unhinged-checkbox")
-        
+                        
         with gr.Row():
             text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
         with gr.Row():
@@ -343,8 +314,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
             seed_output = gr.Textbox(label="Seed Used", elem_id="seed-output")
             final_prompt_output = gr.Textbox(label="Final Prompt", elem_id="final-prompt-output")
         
-        # Update the click function to include the use_unhinged option
-        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev, enhance_prompt_style, enhance_prompt_option, nemo_prompt_style, use_mistral_nemo, use_unhinged], outputs=[image_output, seed_output, final_prompt_output])
-
+        # Adjust the click function to include the API key, use_dev, and enhance_prompt_option as inputs
+        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev, enhance_prompt_style,enhance_prompt_option, enhance_prompt_style, use_mistral_nemo], outputs=[image_output, seed_output, final_prompt_output])
 
 app.launch(show_api=True, share=False)