Update app.py

app.py

@@ -9,6 +9,8 @@ import random
 import numpy as np
 import yaml
 import traceback
+import asyncio
+from transformers import AutoTokenizer, pipeline
 
 try:
     with open("config.yaml", "r") as file:
@@ -24,9 +26,25 @@ DATA_PATH.mkdir(exist_ok=True)
 PREDEFINED_SEED = random.randint(0, MAX_SEED)
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
 
+try:
+    llm_client = pipeline(
+        "text-generation",
+        model="mistralai/Mistral-7B-Instruct-v0.1",
+        device="cpu"
+    )
+except Exception as e:
+    st.error(f"Error al cargar el modelo de lenguaje: {e}")
+    llm_client = None
+
 if not HF_TOKEN_UPSCALER:
     st.warning("HF_TOKEN_UPSCALER no está configurado. Algunas funcionalidades pueden no funcionar.")
 
+def handle_file(file_path):
+    if isinstance(file_path, (str, Path)):
+        with open(file_path, 'rb') as f:
+            return f.read()
+    return file_path
+
 def get_upscale_finegrain(prompt, img_path, upscale_factor):
     try:
         upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
@@ -127,7 +145,7 @@ def save_prompt(prompt):
         st.success("Prompt guardado.")
 
 async def improve_prompt(prompt):
-    if not prompt.strip():
+    if not prompt.strip() or not llm_client:
         return prompt
     
     try:
@@ -142,7 +160,7 @@ async def improve_prompt(prompt):
                 "With this inspiration, realize a cinematic txt2img prompt in English with hyperrealistic elements in 200 characters maximum",
                 "With my idea, make a lifelike and txt2img prompt in English, focusing on photorealistic depth in 200 characters maximum"
             ]
-            response = llm_client.text_generation(
+            result = llm_client(
                 f"{prompt}: {random.choice(instructions)}",
                 max_new_tokens=256,
                 temperature=0.7,
@@ -150,7 +168,8 @@ async def improve_prompt(prompt):
                 repetition_penalty=1.1,
                 do_sample=True
             )
-            return (response['generated_text'] if isinstance(response, dict) else str(response))[:200].strip() or prompt
+            generated_text = result[0]['generated_text'] if isinstance(result, list) else str(result)
+            return generated_text[:200].strip() or prompt
     except Exception as e:
         st.error(f"Error al mejorar prompt: {str(e)}\n{traceback.format_exc()}")
         return prompt
@@ -166,6 +185,8 @@ async def generate_variations(prompt, num_variants, use_enhanced):
             prompts.add(enhanced_prompt)
             if len(prompts) == 1 and num_variants > 1:
                 prompts.add(prompt)
+        if len(prompts) < num_variants:
+            prompts.update([prompt] * (num_variants - len(prompts)))
     except Exception as e:
         st.error(f"Error generando variaciones: {e}")
         return [prompt] * num_variants
@@ -240,7 +261,12 @@ def main():
         prompts = [prompt] * num_variants
 
     if st.sidebar.button("Generar Imagen"):
-        images = gen(prompts, width, height, model_option, num_variants=num_variants)
+        try:
+            images = gen(prompts, width, height, model_option, num_variants=num_variants)
+            if not images:
+                st.error("No se pudieron generar imágenes")
+        except Exception as e:
+            st.error(f"Error durante la generación: {e}")
 
     if generated_image_path and upscale_checkbox:
         upscale_factor = st.sidebar.slider("Factor de Escalado", 1, 4, 2)