Update app.py

app.py

@@ -23,14 +23,22 @@ hf_token = os.getenv("HF_TOKEN")
 
 # Inicializa o modelo base FLUX.1-dev
 base_model = "black-forest-labs/FLUX.1-dev"
+
+# Verifica se CUDA está disponível
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+print(f"✅ Dispositivo detectado: {device}")
+print(f"✅ Usando precisão: {dtype}")
+
 pipe = DiffusionPipeline.from_pretrained(
     base_model, 
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     use_safetensors=True
 )
 
-# Move o modelo para GPU
-pipe.to("cuda")
+# Move o modelo para o dispositivo disponível
+pipe.to(device)
 
 # Definição dos LoRA e Trigger Words
 lora_models = {
@@ -106,7 +114,9 @@ def translate_text(text, source_lang="pt", target_lang="en"):
 def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_option, lora_scale_1, lora_scale_2, cross_attention_scale, auto_translate, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
+    # Cria o gerador no dispositivo correto (cuda ou cpu)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     original_prompt = prompt  # Guarda o prompt original para metadados
     
@@ -150,8 +160,20 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     if lora_option == "Ambos" and cross_attention_scale != 1.0:
         cross_attention_kwargs = {"scale": cross_attention_scale}
 
-    # Gera a imagem com precisão de 16 bits
-    with torch.autocast("cuda"):
+    # Gera a imagem com o dispositivo e precisão adequados
+    if device == "cuda":
+        with torch.autocast("cuda"):
+            image = pipe(
+                prompt=prompt,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                cross_attention_kwargs=cross_attention_kwargs
+            ).images[0]
+    else:
+        # Em CPU não é necessário usar autocast
         image = pipe(
             prompt=prompt,
             num_inference_steps=steps,
@@ -206,7 +228,19 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
 # Interface Gradio
 gr_theme = os.getenv("THEME")
 with gr.Blocks(theme=gr_theme) as app:
-    gr.Markdown("# Paula & Vivi Image Generator")
+    device_info = "GPU" if torch.cuda.is_available() else "CPU"
+    gr.Markdown(f"# Paula & Vivi Image Generator (Rodando em {device_info})")
+    
+    if device == "cpu":
+        gr.Markdown("""
+        ⚠️ **Aviso: Executando em CPU**
+        
+        Este aplicativo está rodando em CPU, o que significa que a geração de imagens será significativamente mais lenta.
+        Considere usar um ambiente com GPU para melhor performance.
+        """)
+    else:
+        gpu_info = torch.cuda.get_device_name(0)
+        gr.Markdown(f"✅ **GPU Detectada**: {gpu_info}")
     
     with gr.Row():
         with gr.Column(scale=2):
@@ -215,9 +249,14 @@ with gr.Blocks(theme=gr_theme) as app:
             
             with gr.Accordion("Configurações Básicas", open=True):
                 cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
-                steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=32)
+                steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=32 if device == "cuda" else 20)
                 width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=768)  
-                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=1024)  
+                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=1024)
+            
+            # Definir valores padrão menores se estiver em CPU
+            if device == "cpu":
+                width.value = 512
+                height.value = 512
                 randomize_seed = gr.Checkbox(False, label="Randomize seed")
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=556215326)