Update app.py

app.py

@@ -13,39 +13,57 @@ from datetime import datetime
 
 import sqlite3
 from datetime import datetime
-
+import sqlite3
+from datetime import datetime
 
 # Initialize the database
 def init_db(file='logs.db'):
     conn = sqlite3.connect(file)
     c = conn.cursor()
-    c.execute('''CREATE TABLE IF NOT EXISTS logs
-                 (timestamp TEXT, message TEXT)''')
+    c.execute('''CREATE TABLE IF NOT EXISTS logs (
+                 id INTEGER PRIMARY KEY AUTOINCREMENT,
+                 timestamp TEXT,
+                 prompt TEXT,
+                 is_negative INTEGER,
+                 steps INTEGER,
+                 cfg_scale REAL,
+                 sampler TEXT,
+                 seed INTEGER,
+                 strength REAL,
+                 use_dev INTEGER,
+                 enhance_prompt_style TEXT,
+                 enhance_prompt_option INTEGER,
+                 nemo_enhance_prompt_style TEXT,
+                 use_mistral_nemo INTEGER
+                 )''')
     conn.commit()
     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):
-    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}"
-    if huggingface_api_key:
-        log_message += f"huggingface_api_key='{huggingface_api_key}'"
-    
-    conn = sqlite3.connect('acces_log.log')
+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):
+    conn = sqlite3.connect('logs.db')
     c = conn.cursor()
-    c.execute("INSERT INTO logs VALUES (?, ?)", (datetime.now().isoformat(), log_message))
+    c.execute('''INSERT INTO logs (timestamp, prompt, is_negative, steps, cfg_scale, sampler, seed, strength, use_dev, 
+                                   enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo) 
+                 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
+              (datetime.now().isoformat(), prompt, is_negative, steps, cfg_scale, sampler, seed, strength, 
+               use_dev, enhance_prompt_style, enhance_prompt_option, nemo_enhance_prompt_style, use_mistral_nemo))
+    
     conn.commit()
     conn.close()
-    
+
+
+
 # os.makedirs('assets', exist_ok=True)
 if not os.path.exists('icon.png'):
     os.system("wget -O icon.png https://huggingface.co/spaces/K00B404/FLUX.1-Dev-Serverless-darn-enhanced-prompt/resolve/main/edge.png")
 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"
 timeout = 100
-init_db('acces_log.log')
+
+
+init_db()
 
 # Set up logging
 logging.basicConfig(filename='access.log', level=logging.INFO,
@@ -53,15 +71,6 @@ logging.basicConfig(filename='access.log', level=logging.INFO,
 
 
 
-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}, "
-    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}"
-    if huggingface_api_key:
-        log_message += f"huggingface_api_key='{huggingface_api_key}'"
-    logging.info(log_message)
-
 def check_ubuse(prompt,word_list=["little girl"]):
     for word in word_list:
         if word in prompt:
@@ -141,33 +150,36 @@ 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):
-    
-    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
+
+
+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 the request (WITHOUT storing API keys)
+    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)
+
+    # Determine 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
+    # API key handling
     is_api_call = huggingface_api_key is not None
-
     if is_api_call:
-        # Use the environment variable for the API key in GUI mode
-        API_TOKEN = os.getenv("HF_READ_TOKEN")
+        API_TOKEN = os.getenv("HF_READ_TOKEN")  # Use env var in GUI mode
     else:
-        # Validate the API key if it's an API call
         if huggingface_api_key == "":
             raise gr.Error("API key is required for API calls.")
         API_TOKEN = huggingface_api_key
     
     headers = {"Authorization": f"Bearer {API_TOKEN}"} 
 
-    if prompt == "" or prompt is None:
+    if not prompt:  # Simplified check
         return None, None, None
 
     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
@@ -181,7 +193,7 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     final_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'\033[1mGeneration {key}:\033[0m {final_prompt}')
 
-    # If seed is -1, generate a random seed and use it
+    # Ensure seed is always positive
     if seed == -1:
         seed = random.randint(1, 1000000000)
 
@@ -207,7 +219,7 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
         image = Image.open(io.BytesIO(image_bytes))
         print(f'\033[1mGeneration {key} completed!\033[0m ({final_prompt})')
 
-        # Save the image to a file and return the file path and seed
+        # Save the image and return path
         output_path = f"./output_{key}.png"
         image.save(output_path)
         
@@ -215,8 +227,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
-
-
   
 title_html="""
     <center>
@@ -431,7 +441,9 @@ with gr.Blocks( 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")
         
+        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]
+        
         # 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])
+        text_button.click(query, inputs=inputs, outputs=[image_output, seed_output, final_prompt_output])
 
 app.launch(show_api=True, share=False)