Update app.py

app.py

@@ -5,13 +5,19 @@ import os, random, numpy as np, yaml, time
 from dataclasses import dataclass
 from typing import List
 from huggingface_hub import InferenceClient
+import google.generativeai as genai
 
 st.set_page_config(layout="wide")
 HF_TOKEN = os.getenv("HF_TOKEN")
+GEMINI_TOKEN = os.getenv("GEMINI_TOKEN")
 
 if not HF_TOKEN:
-    st.error("Error en el token! 'HF_TOKEN'.")
-    st.stop() 
+    st.error("Error en el token de Hugging Face! 'HF_TOKEN'.")
+    st.stop()
+
+if not GEMINI_TOKEN:
+    st.error("Error en el token de Gemini! 'GEMINI_TOKEN'.")
+    st.stop()
 
 try:
     with open("config.yaml", "r") as file: 
@@ -29,10 +35,15 @@ MAX_SEED = AppConfig.MAX_SEED
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
-def get_inference_client():
+def get_hf_client():
     return InferenceClient(token=HF_TOKEN)
 
-client = get_inference_client()
+def get_gemini_client():
+    genai.configure(api_key=GEMINI_TOKEN)
+    return genai.GenerativeModel('gemini-1.5-flash-8b')
+
+hf_client = get_hf_client()
+gemini_client = get_gemini_client()
 
 def authenticate_user(username, password): 
     return username == credentials["username"] and password == credentials["password"]
@@ -40,7 +51,7 @@ def authenticate_user(username, password):
 def list_saved_images(): 
     return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)
 
-def enhance_prompt(text, client=client, use_enhancement=True):
+def enhance_prompt(text, client=None, use_enhancement=True):
     if not use_enhancement:
         return text[:200]
     
@@ -48,15 +59,21 @@ def enhance_prompt(text, client=client, use_enhancement=True):
         return text[:200]
     
     try:
-        enhanced = client.text_generation(
-            "Generate a photorealistic, detailed txt2img prompt: " + text,
-            model="mistralai/Mixtral-8x7B-Instruct-v0.1",)
-        return enhanced[:200]
+        enhancement_instruction = "Enhance this text description to be more suitable for text-to-image generation. Focus on vivid, descriptive language that will help an AI generate a photorealistic image. Be specific about colors, composition, lighting, and key details."
+        
+        response = client.generate_content(f"{enhancement_instruction}\n\nOriginal prompt: {text}")
+        enhanced_prompt = response.text.strip()
         
+        return enhanced_prompt[:200]
+    
     except Exception as e:
         st.warning(f"Prompt enhancement error: {e}")
         return text[:200]
 
+def save_prompt(image_name, enhanced_prompt):
+    with open(DATA_PATH / "prompts.txt", "a") as f: 
+        f.write(f"{image_name}: {enhanced_prompt}\n")
+
 def generate_variations(prompt, num_variants=8, use_enhanced=True):
     instructions = [
         "Photorealistic description for txt2img prompt: ",
@@ -69,12 +86,12 @@ def generate_variations(prompt, num_variants=8, use_enhanced=True):
         "Lifelike txt2img, focusing on photorealistic depth: "
     ]
     if use_enhanced:
-        prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}") for i in range(num_variants)]
+        prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}", client=gemini_client, use_enhancement=True) for i in range(num_variants)]
     else: 
-        prompts = [prompt] * num_variants
+        prompts = [enhance_prompt(prompt, use_enhancement=False) for i in range(num_variants)]
     return prompts
 
-def generate_image(prompt, width, height, seed, model_name, client=client):
+def generate_image(prompt, width, height, seed, model_name, client=None):
     if not client:
         st.error("No Hugging Face client available")
         return None, seed, None
@@ -82,7 +99,7 @@ def generate_image(prompt, width, height, seed, model_name, client=client):
     try:
         with st.spinner("Generando imagen..."):
             seed = int(seed) if seed != -1 else random.randint(0, AppConfig.MAX_SEED)
-            enhanced_prompt = enhance_prompt(prompt)
+            enhanced_prompt = enhance_prompt(prompt, client=gemini_client)
             image = client.text_to_image(
                 prompt=enhanced_prompt, 
                 height=height, 
@@ -106,7 +123,7 @@ def gen(prompts, width, height, model_name, num_variants=8):
     for i in range(num_variants):
         current_prompt = prompts[i] if len(prompts) > i else prompts[-1]
         with st.spinner(f"Generando imagen {i+1}/{num_variants}"):
-            image, used_seed, enhanced_prompt = generate_image(current_prompt, width, height, seeds[i], model_name)
+            image, used_seed, enhanced_prompt = generate_image(current_prompt, width, height, seeds[i], model_name, client=hf_client)
             if image:
                 image_path = DATA_PATH / f"generated_image_{used_seed}.jpg"
                 image.save(image_path)
@@ -173,8 +190,8 @@ def main():
     
     st.title("Flux +Upscale +Prompt Enhancer")
     
-    if not client:
-        st.error("No se pudo establecer conexión con Hugging Face. Verifique sus tokens.")
+    if not hf_client or not gemini_client:
+        st.error("No se pudo establecer conexión con los servicios. Verifique sus tokens.")
         return
     
     prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)