Update app.py

app.py

@@ -2,114 +2,98 @@ import gradio as gr
 import torch
 import spaces
 import os
+import sys
 import shutil
-import glob
-from diffusers import DiffusionPipeline
+import importlib.util
 from huggingface_hub import snapshot_download
 
 # -----------------------------------------------------------------------------
-# 1. FUNCIÓN DE DIAGNÓSTICO Y REPARACIÓN PROFUNDA
+# CONFIGURACIÓN
 # -----------------------------------------------------------------------------
 MODEL_ID = "NewBie-AI/NewBie-image-Exp0.1"
-LOCAL_DIR = "./newbie_fixed_model"
+GITHUB_REPO_URL = "https://github.com/NewBie-AI/NewBie" # El origen del código perdido
+LOCAL_MODEL_DIR = "./model_weights"
+LOCAL_CODE_DIR = "./newbie_code"
 
-def list_all_files(directory):
-    """Función auxiliar para ver qué rayos se descargó"""
-    print(f"\n📂 LISTADO DE ARCHIVOS EN {directory}:")
-    for root, dirs, files in os.walk(directory):
-        level = root.replace(directory, '').count(os.sep)
-        indent = ' ' * 4 * (level)
-        print(f"{indent}{os.path.basename(root)}/")
-        subindent = ' ' * 4 * (level + 1)
-        for f in files:
-            print(f"{subindent}{f}")
-    print("------------------------------------------------\n")
-
-def load_deep_fixed_pipeline():
-    print(f"🛠️ Iniciando protocolo de Búsqueda Profunda para {MODEL_ID}...")
+# -----------------------------------------------------------------------------
+# FUNCIÓN DE RESCATE: CLONAR CÓDIGO + DESCARGAR PESOS
+# -----------------------------------------------------------------------------
+def load_hybrid_pipeline():
+    print(f"🚨 INICIANDO PROTOCOLO DE RESCATE PARA {MODEL_ID}...")
     
-    # 1. Descargar repositorio completo
-    if not os.path.exists(LOCAL_DIR):
-        print("   ⬇️ Descargando snapshot...")
-        snapshot_download(repo_id=MODEL_ID, local_dir=LOCAL_DIR)
+    # 1. Descargar Pesos (Hugging Face)
+    if not os.path.exists(LOCAL_MODEL_DIR):
+        print("   ⬇️ Descargando pesos del modelo (Safetensors)...")
+        snapshot_download(
+            repo_id=MODEL_ID, 
+            local_dir=LOCAL_MODEL_DIR,
+            ignore_patterns=["*.msgpack", "*.bin"] # Optimizamos descarga
+        )
+
+    # 2. Descargar Código (GitHub)
+    if not os.path.exists(LOCAL_CODE_DIR):
+        print(f"   ⬇️ Clonando código fuente desde {GITHUB_REPO_URL}...")
+        # Usamos git clone para traer el código que falta en HF
+        os.system(f"git clone {GITHUB_REPO_URL} {LOCAL_CODE_DIR}")
     
-    # 2. IMPRIMIR ESTRUCTURA (Para depuración si falla)
-    list_all_files(LOCAL_DIR)
-
-    # 3. BUSCAR ARCHIVOS DE CÓDIGO (.py) EN CUALQUIER SUBDIRECTORIO
-    all_py_files = []
-    for root, dirs, files in os.walk(LOCAL_DIR):
+    # 3. Preparar el entorno de Python
+    # Añadimos la carpeta clonada al path para que Python "vea" los archivos nuevos
+    sys.path.append(os.path.abspath(LOCAL_CODE_DIR))
+    
+    # 4. BUSCAR LA CLASE 'NewbiePipeline' MANUALMENTE
+    print("   🕵️‍♂️ Buscando la clase perdida 'NewbiePipeline' en el código clonado...")
+    pipeline_class = None
+    
+    # Escaneamos recursivamente el repo de GitHub clonado
+    for root, dirs, files in os.walk(LOCAL_CODE_DIR):
         for file in files:
             if file.endswith(".py"):
-                full_path = os.path.join(root, file)
-                all_py_files.append(full_path)
-    
-    print(f"   🔎 Se encontraron {len(all_py_files)} archivos Python: {all_py_files}")
-
-    if not all_py_files:
-        raise RuntimeError("❌ ERROR FATAL: No se encontró NINGÚN archivo .py en el repositorio. El modelo no se puede ejecutar.")
-
-    # 4. PREPARAR CARPETA TRANSFORMER (Donde el config busca el código)
-    transformer_folder = os.path.join(LOCAL_DIR, "transformer")
-    os.makedirs(transformer_folder, exist_ok=True)
+                path = os.path.join(root, file)
+                try:
+                    with open(path, "r", encoding="utf-8", errors="ignore") as f:
+                        if "class NewbiePipeline" in f.read():
+                            print(f"   🎯 ¡CÓDIGO ENCONTRADO EN!: {file}")
+                            
+                            # Importación dinámica (Magia negra de Python)
+                            spec = importlib.util.spec_from_file_location("dynamic_pipeline", path)
+                            module = importlib.util.module_from_spec(spec)
+                            sys.modules["dynamic_pipeline"] = module
+                            spec.loader.exec_module(module)
+                            pipeline_class = getattr(module, "NewbiePipeline")
+                            break
+                except Exception:
+                    continue
+        if pipeline_class: break
     
-    # Crear __init__.py vacío para que sea un módulo
-    with open(os.path.join(transformer_folder, "__init__.py"), "w") as f: f.write("")
-
-    # 5. ESTRATEGIA DE COPIA INTELIGENTE
-    # Copiamos TODOS los archivos .py encontrados a la carpeta 'transformer/'
-    # y también a la raíz, para asegurar que se encuentren.
-    for py_file in all_py_files:
-        filename = os.path.basename(py_file)
-        
-        # Copiar a transformer/
-        shutil.copy(py_file, os.path.join(transformer_folder, filename))
-        
-        # Copiar a raíz (si no está ya ahí)
-        root_dest = os.path.join(LOCAL_DIR, filename)
-        if not os.path.exists(root_dest):
-            shutil.copy(py_file, root_dest)
-
-        # DETECTAR EL CANDIDATO PRINCIPAL PARA 'transformer.py'
-        # Buscamos archivos que contengan "Transformer" o "Model" en su contenido
-        # O que se llamen 'modeling_...'
-        if "modeling" in filename or "transformer" in filename.lower():
-            print(f"   ✅ Posible archivo de modelado detectado: {filename}")
-            # Lo forzamos como transformer.py
-            shutil.copy(py_file, os.path.join(transformer_folder, "transformer.py"))
+    if not pipeline_class:
+        raise RuntimeError("❌ No se encontró 'class NewbiePipeline' ni siquiera en el GitHub. El código ha cambiado.")
 
-    # Si después de todo no existe transformer.py, usamos el archivo más grande encontrado
-    if not os.path.exists(os.path.join(transformer_folder, "transformer.py")):
-        print("   ⚠️ No se detectó un nombre obvio. Usando el archivo .py más grande como transformer.py")
-        largest_py = max(all_py_files, key=os.path.getsize)
-        shutil.copy(largest_py, os.path.join(transformer_folder, "transformer.py"))
-
-    # 6. CARGAR
-    print("   🚀 Intentando cargar pipeline...")
-    # Usamos local_files_only para obligarlo a usar nuestra estructura hackeada
-    pipe = DiffusionPipeline.from_pretrained(
-        LOCAL_DIR,
+    # 5. INSTANCIAR EL PIPELINE
+    print("   🚀 Conectando código clonado con pesos descargados...")
+    pipe = pipeline_class.from_pretrained(
+        LOCAL_MODEL_DIR,
         torch_dtype=torch.bfloat16,
         trust_remote_code=True,
         local_files_only=True
     )
+    
     return pipe
 
-# Carga inicial protegida
+# Ejecutar carga
 pipe = None
 try:
-    pipe = load_deep_fixed_pipeline()
+    pipe = load_hybrid_pipeline()
+    print("   ✅ ¡MODELO CARGADO EXITOSAMENTE!")
 except Exception as e:
-    print(f"❌ Error durante la carga: {e}")
-    # No detenemos el script para que Gradio pueda mostrar el error en pantalla si es necesario
+    print(f"❌ ERROR CRÍTICO: {e}")
 
 # -----------------------------------------------------------------------------
-# 2. LÓGICA ZEROGPU
+# LÓGICA ZEROGPU
 # -----------------------------------------------------------------------------
 @spaces.GPU(duration=120)
 def generate_image(prompt, negative_prompt, steps, cfg, width, height):
     if pipe is None:
-        raise gr.Error("El modelo falló al cargar. Revisa los logs de la consola (Files listed above).")
+        raise gr.Error("El modelo no está cargado. Revisa la consola.")
     
     print("🎨 Generando...")
     pipe.to("cuda")
@@ -125,13 +109,15 @@ def generate_image(prompt, negative_prompt, steps, cfg, width, height):
         ).images[0]
         return image
     except Exception as e:
-        raise gr.Error(f"Error de inferencia: {e}")
+        raise gr.Error(f"Error generando imagen: {e}")
 
 # -----------------------------------------------------------------------------
-# 3. INTERFAZ (SIN ARGUMENTO THEME)
+# INTERFAZ
 # -----------------------------------------------------------------------------
 css = """
+<style>
 .container { max-width: 900px; margin: auto; }
+</style>
 """
 
 DEFAULT_PROMPT = """<character_1>
@@ -141,27 +127,22 @@ DEFAULT_PROMPT = """<character_1>
 <action>standing, holding_fan</action>
 </character_1>
 <general_tags>
-<quality>best quality, masterpiece, 4k</quality>
 <style>anime, vivid_colors</style>
 </general_tags>"""
 
-DEFAULT_NEG = "low quality, bad anatomy, worst quality, watermark, text"
-
-# ELIMINADO theme=... y css=... del constructor para compatibilidad total
 with gr.Blocks() as demo:
-    gr.HTML(f"<style>{css}</style>") # CSS inyectado manualmente
-    gr.Markdown("# ⛩️ NewBie Anime Generator (Deep Fix)")
+    gr.HTML(css)
+    gr.Markdown("# ⛩️ NewBie Anime (GitHub Rescue Edition)")
     
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt (XML)", value=DEFAULT_PROMPT, lines=10)
-            neg = gr.Textbox(label="Negative", value=DEFAULT_NEG)
+            prompt = gr.Textbox(label="Prompt (XML)", value=DEFAULT_PROMPT, lines=8)
+            neg = gr.Textbox(label="Negative", value="low quality, bad anatomy")
             btn = gr.Button("Generar", variant="primary")
-            with gr.Row():
-                steps = gr.Slider(10, 50, value=28, label="Pasos")
-                cfg = gr.Slider(1, 15, value=7.0, label="CFG")
-                width = gr.Slider(512, 1280, value=1024, step=64, label="Ancho")
-                height = gr.Slider(512, 1280, value=1024, step=64, label="Alto")
+            steps = gr.Slider(10, 50, value=28, label="Pasos")
+            cfg = gr.Slider(1, 15, value=7.0, label="CFG")
+            width = gr.Slider(512, 1280, value=1024, step=64, label="Ancho")
+            height = gr.Slider(512, 1280, value=1024, step=64, label="Alto")
         with gr.Column():
             out = gr.Image(label="Resultado")