Update app.py

app.py

@@ -41,8 +41,11 @@ def log_request(prompt, is_negative, steps, cfg_scale, sampler, seed, strength,
 # 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')
 
@@ -51,6 +54,14 @@ 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}, "
@@ -141,11 +152,17 @@ 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):
+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)
     
-    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
-    api_url = API_URL_DEV if use_dev else API_URL
+    if use_unhinged:
+        api_url = API_URL_UNHINGED
+    elif use_dev:
+        api_url = API_URL_DEV
+    else:
+        api_url = 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
@@ -166,7 +183,6 @@ 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
@@ -184,6 +200,17 @@ 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,
@@ -214,6 +241,7 @@ 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 {
@@ -276,7 +304,6 @@ css1 = """
     font-weight: bold;
 }
 """
-
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
     gr.HTML("""
         <center>
@@ -305,7 +332,8 @@ 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():
@@ -314,7 +342,8 @@ 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")
         
-        # 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])
+        # 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])
+
 
 app.launch(show_api=True, share=False)