Update app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-# Configuración para usar bfloat16 y CUDA si está disponible
-dtype = torch.bfloat16
+# Configuración para usar bfloat16 si está disponible y detección automática del dispositivo
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+# Detectar si CUDA está disponible, en caso contrario usar CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Inicializar el modelo solo una vez y cargarlo en RAM y GPU
+# Inicializar el modelo solo una vez y cargarlo en RAM y GPU/CPU
 pipe = None
 
 def load_model():
@@ -21,11 +23,11 @@ def load_model():
             "black-forest-labs/FLUX.1-schnell", 
             torch_dtype=dtype
         )
-        # Despachar los pesos a la GPU, evitando acceder a named_parameters
+        # Despachar los pesos al dispositivo adecuado (GPU o CPU)
         pipe = load_checkpoint_and_dispatch(
             pipe, 
             "black-forest-labs/FLUX.1-schnell", 
-            device_map="auto",  # Automatiza el uso de RAM y GPU
+            device_map="auto",  # Automatiza el uso de RAM, GPU o CPU
             offload_folder=None  # Evita que se almacenen los pesos temporalmente en el disco
         )
         pipe.to(device)
@@ -81,7 +83,8 @@ with gr.Blocks(css=css) as demo:
                 show_label=False,
                 max_lines=1,
                 placeholder="Enter your prompt",
-                container=False 
+                container=False,
+                max_length=10000000  # Máximo de 10,000,000 caracteres
             )
             
             run_button = gr.Button("Run", scale=0)