Update app.py

app.py

@@ -3,123 +3,138 @@ import torch
 import spaces
 from diffusers import DiffusionPipeline
 import os
+import shutil
+import json
+from huggingface_hub import snapshot_download
 
 # -----------------------------------------------------------------------------
-# 1. CONFIGURACIÓN Y CARGA DEL MODELO
+# 1. FUNCIÓN DE REPARACIÓN Y CARGA (La solución al error)
 # -----------------------------------------------------------------------------
 MODEL_ID = "NewBie-AI/NewBie-image-Exp0.1"
+LOCAL_DIR = "./newbie_fixed_model"
 
-print(f"🔄 Iniciando carga del modelo: {MODEL_ID}")
-print("   Nota: Esto puede tardar unos minutos la primera vez mientras se descargan los pesos.")
+def load_fixed_pipeline():
+    print(f"🛠️ Iniciando protocolo de reparación para {MODEL_ID}...")
+    
+    # 1. Descargar el repositorio completo a una carpeta local
+    if not os.path.exists(LOCAL_DIR):
+        print(f"   Descargando snapshot del modelo... (Esto tardará un poco)")
+        snapshot_download(
+            repo_id=MODEL_ID,
+            local_dir=LOCAL_DIR,
+            ignore_patterns=["*.msgpack", "*.bin"] # Ignoramos binarios redundantes si hay safetensors
+        )
+    
+    # 2. Analizar y corregir el fallo de estructura (transformer/transformer.py)
+    transformer_folder = os.path.join(LOCAL_DIR, "transformer")
+    
+    # Si la carpeta 'transformer' no existe, la creamos
+    if not os.path.exists(transformer_folder):
+        os.makedirs(transformer_folder, exist_ok=True)
+        print("   Carpeta 'transformer' creada.")
+        
+        # BUSCAMOS EL ARCHIVO PERDIDO
+        # El repo suele tener el archivo en la raíz con nombres como 'transformer.py' o 'modeling_newbie.py'
+        # Vamos a buscar archivos python candidatos en la raíz
+        candidates = [f for f in os.listdir(LOCAL_DIR) if f.endswith(".py") and "test" not in f]
+        
+        # Buscamos específicamente uno que parezca el del transformer
+        target_file = None
+        for f in candidates:
+            if "transformer" in f or "modeling" in f:
+                target_file = f
+                break
+        
+        # Si encontramos el archivo, lo copiamos a la carpeta transformer/transformer.py
+        if target_file:
+            src = os.path.join(LOCAL_DIR, target_file)
+            dst = os.path.join(transformer_folder, "transformer.py")
+            shutil.copy(src, dst)
+            print(f"   ✅ Archivo reparado: Movido {target_file} a transformer/transformer.py")
+            
+            # También necesitamos un __init__.py vacío para que sea un módulo
+            with open(os.path.join(transformer_folder, "__init__.py"), "w") as f:
+                f.write("")
+        else:
+            print("   ⚠️ ADVERTENCIA: No se encontró un archivo Python candidato obvio en la raíz.")
 
-# Cargamos el pipeline. 
-# SOLUCIÓN APLICADA: 'custom_pipeline' fuerza la descarga del script de gestión del repo,
-# evitando el error de ruta 'transformer/transformer.py' que te salía antes.
-pipe = DiffusionPipeline.from_pretrained(
-    MODEL_ID,
-    custom_pipeline=MODEL_ID,     # <--- EL CAMBIO CLAVE
-    torch_dtype=torch.bfloat16,   # bfloat16 es nativo y más rápido en A100 (ZeroGPU)
-    trust_remote_code=True
-)
+    # 3. Cargar desde la carpeta local reparada
+    print("   Cargando pipeline desde disco local...")
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            LOCAL_DIR,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=True
+        )
+        return pipe
+    except Exception as e:
+        print(f"   ❌ Error fatal cargando localmente: {e}")
+        # Intento desesperado: Carga remota estándar si el parche falla, 
+        # pero quitando custom_pipeline que a veces causa bucles
+        return DiffusionPipeline.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True
+        )
+
+# Ejecutamos la carga al iniciar la app
+pipe = load_fixed_pipeline()
 
 # -----------------------------------------------------------------------------
-# 2. FUNCIÓN DE GENERACIÓN (ZeroGPU)
+# 2. LÓGICA ZEROGPU
 # -----------------------------------------------------------------------------
-# El decorador @spaces.GPU maneja la asignación de hardware.
-# duration=120 da tiempo suficiente para imágenes grandes sin timeout.
 @spaces.GPU(duration=120)
 def generate_image(prompt, negative_prompt, steps, cfg, width, height):
-    print("🚀 ZeroGPU Asignada. Moviendo modelo a CUDA y generando...")
-    
-    # Mover a GPU solo dentro de la función decorada
+    print("🚀 ZeroGPU Asignada. Generando...")
     pipe.to("cuda")
     
-    # Ejecutar inferencia
-    try:
-        image = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=int(steps),
-            guidance_scale=float(cfg),
-            width=int(width),
-            height=int(height)
-        ).images[0]
-        return image
-    except Exception as e:
-        print(f"Error durante la generación: {e}")
-        return None
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=int(steps),
+        guidance_scale=float(cfg),
+        width=int(width),
+        height=int(height)
+    ).images[0]
+    return image
 
 # -----------------------------------------------------------------------------
-# 3. INTERFAZ GRÁFICA (GRADIO)
+# 3. INTERFAZ GRADIO
 # -----------------------------------------------------------------------------
-# CSS para mejorar un poco la estética
 css = """
 .container { max-width: 900px; margin: auto; }
-textarea { font-family: monospace; }
 """
 
-# Prompt por defecto optimizado para este modelo (Formato XML)
 DEFAULT_PROMPT = """<character_1>
 <gender>1girl</gender>
-<appearance>red_eyes, white_hair, long_hair, floating_hair</appearance>
-<clothing>japanese_clothes, kimono, floral_print</clothing>
-<action>standing, holding_fan, looking_at_viewer</action>
+<appearance>blue_eyes, silver_hair, long_hair</appearance>
+<clothing>school_uniform, serafuku</clothing>
+<action>standing, smile</action>
 </character_1>
 <general_tags>
-<quality>best quality, masterpiece, 4k, highres</quality>
-<style>anime, vivid_colors, cherry_blossoms</style>
+<quality>masterpiece, best quality, 4k</quality>
+<style>anime, vivid_colors</style>
 </general_tags>"""
 
-DEFAULT_NEGATIVE = "low quality, bad anatomy, worst quality, watermark, text, error, jpeg artifacts, signature"
+DEFAULT_NEGATIVE = "low quality, bad anatomy, worst quality, watermark, text"
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # ⛩️ NewBie Anime Generator (ZeroGPU Edition)
-    
-    Este espacio utiliza el modelo experimental **NewBie-image-Exp0.1**. 
-    Este modelo entiende mejor los prompts si usas una estructura **XML** (ver ejemplo abajo).
-    """)
+    gr.Markdown("# ⛩️ NewBie Anime Generator (Auto-Fixed Edition)")
+    gr.Markdown("Este espacio incluye un script de **auto-reparación** para solucionar el error de estructura del repositorio NewBie-Exp0.1.")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt_input = gr.Textbox(
-                label="Prompt (Estructura XML Recomendada)", 
-                value=DEFAULT_PROMPT, 
-                lines=12,
-                placeholder="Escribe tu prompt aquí..."
-            )
-            neg_prompt_input = gr.Textbox(
-                label="Negative Prompt", 
-                value=DEFAULT_NEGATIVE, 
-                lines=2
-            )
-            
-            with gr.Accordion("⚙️ Configuración Avanzada", open=False):
-                with gr.Row():
-                    width_slider = gr.Slider(512, 1280, value=1024, step=64, label="Ancho")
-                    height_slider = gr.Slider(512, 1280, value=1024, step=64, label="Alto")
-                steps_slider = gr.Slider(10, 50, value=28, step=1, label="Pasos (Steps)")
-                cfg_slider = gr.Slider(1, 15, value=7.0, step=0.1, label="Guidance Scale")
-            
-            btn_run = gr.Button("🎨 Generar Imagen", variant="primary", scale=1)
-        
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Resultado", type="pil", interactive=False)
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt (XML)", value=DEFAULT_PROMPT, lines=10)
+            neg_prompt_input = gr.Textbox(label="Negative", value=DEFAULT_NEGATIVE)
+            with gr.Row():
+                steps = gr.Slider(10, 50, value=28, label="Pasos")
+                cfg = gr.Slider(1, 15, value=7.0, label="CFG")
+            btn = gr.Button("Generar", variant="primary")
+        with gr.Column():
+            out = gr.Image(label="Resultado")
 
-    # Evento de clic
-    btn_run.click(
-        fn=generate_image,
-        inputs=[
-            prompt_input, 
-            neg_prompt_input, 
-            steps_slider, 
-            cfg_slider, 
-            width_slider, 
-            height_slider
-        ],
-        outputs=output_image
-    )
+    btn.click(generate_image, inputs=[prompt_input, neg_prompt_input, steps, cfg], outputs=out)
 
-# Lanzar la aplicación
 if __name__ == "__main__":
     demo.launch()