Update app.py

app.py

@@ -1,92 +1,115 @@
 import gradio as gr
 import torch
 import spaces
-from diffusers import DiffusionPipeline
 import os
 import shutil
-import json
+import sys
+import importlib.util
 from huggingface_hub import snapshot_download
 
 # -----------------------------------------------------------------------------
-# 1. FUNCIÓN DE REPARACIÓN Y CARGA (La solución al error)
+# 1. FUNCIÓN DE INGENIERÍA INVERSA Y CARGA
 # -----------------------------------------------------------------------------
 MODEL_ID = "NewBie-AI/NewBie-image-Exp0.1"
 LOCAL_DIR = "./newbie_fixed_model"
 
-def load_fixed_pipeline():
-    print(f"🛠️ Iniciando protocolo de reparación para {MODEL_ID}...")
+def load_hard_fixed_pipeline():
+    print(f"🛠️ Iniciando protocolo de carga manual para {MODEL_ID}...")
     
-    # 1. Descargar el repositorio completo a una carpeta local
+    # --- PASO 1: Descarga ---
     if not os.path.exists(LOCAL_DIR):
-        print(f"   Descargando snapshot del modelo... (Esto tardará un poco)")
+        print(f"   Descargando snapshot del modelo...")
         snapshot_download(
             repo_id=MODEL_ID,
             local_dir=LOCAL_DIR,
-            ignore_patterns=["*.msgpack", "*.bin"] # Ignoramos binarios redundantes si hay safetensors
+            ignore_patterns=["*.msgpack", "*.bin"]
         )
     
-    # 2. Analizar y corregir el fallo de estructura (transformer/transformer.py)
+    # --- PASO 2: Arreglar estructura de carpetas (Transformer) ---
     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
+        # Buscar el archivo del transformer
         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")
+            # Movemos cualquier cosa que parezca transformer o modeling
+            if "transformer" in f.lower() or "modeling" in f.lower():
+                src = os.path.join(LOCAL_DIR, f)
+                # El model_index suele buscar 'transformer.py' o 'modeling_transformer.py'
+                # Lo renombramos a transformer.py para estandarizar
+                dst = os.path.join(transformer_folder, "transformer.py") 
+                if not os.path.exists(dst): # Solo si no existe ya
+                    shutil.copy(src, dst)
+                    print(f"   📂 Archivo movido: {f} -> transformer/transformer.py")
+    
+    # Crear __init__.py en transformer para que sea importable
+    init_path = os.path.join(transformer_folder, "__init__.py")
+    if not os.path.exists(init_path):
+        with open(init_path, "w") as f: f.write("")
+
+    # --- PASO 3: Importación Dinámica del Pipeline (El Fix Crítico) ---
+    # Añadimos el directorio al path para que Python encuentre los módulos internos
+    sys.path.append(os.path.abspath(LOCAL_DIR))
+    
+    pipeline_class = None
+    
+    # Buscamos qué archivo contiene la clase "NewbiePipeline"
+    py_files = [f for f in os.listdir(LOCAL_DIR) if f.endswith(".py")]
+    
+    for py_file in py_files:
+        try:
+            file_path = os.path.join(LOCAL_DIR, py_file)
+            with open(file_path, "r", encoding="utf-8") as f:
+                content = f.read()
             
-            # 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.")
+            if "class NewbiePipeline" in content:
+                print(f"   🎯 Clase encontrada en: {py_file}")
+                
+                # Importar el módulo manualmente
+                spec = importlib.util.spec_from_file_location("newbie_pipeline_module", file_path)
+                module = importlib.util.module_from_spec(spec)
+                sys.modules["newbie_pipeline_module"] = module
+                spec.loader.exec_module(module)
+                
+                # Obtener la clase
+                pipeline_class = getattr(module, "NewbiePipeline")
+                break
+        except Exception as e:
+            print(f"   Saltando archivo {py_file}: {e}")
+            continue
 
-    # 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
-        )
+    if pipeline_class is None:
+        raise ValueError("❌ No se encontró la clase 'NewbiePipeline' en ningún archivo .py del repositorio.")
+
+    # --- PASO 4: Carga del Modelo usando la Clase Importada ---
+    print("   🚀 Cargando pesos usando la clase personalizada...")
+    pipe = pipeline_class.from_pretrained(
+        LOCAL_DIR,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        # Importante: Como ya estamos usando la clase directa, no necesitamos custom_pipeline
+        # pero mantenemos trust_remote_code por si acaso el transformer lo requiere
+        trust_remote_code=True 
+    )
+    
+    return pipe
 
-# Ejecutamos la carga al iniciar la app
-pipe = load_fixed_pipeline()
+# Ejecutar carga
+try:
+    pipe = load_hard_fixed_pipeline()
+except Exception as e:
+    print(f"❌ ERROR CRÍTICO EN CARGA: {e}")
+    pipe = None
 
 # -----------------------------------------------------------------------------
 # 2. LÓGICA ZEROGPU
 # -----------------------------------------------------------------------------
 @spaces.GPU(duration=120)
 def generate_image(prompt, negative_prompt, steps, cfg, width, height):
-    print("🚀 ZeroGPU Asignada. Generando...")
+    if pipe is None:
+        raise gr.Error("El modelo no se pudo cargar. Revisa los logs.")
+        
+    print("🎨 Generando...")
     pipe.to("cuda")
     
     image = pipe(
@@ -100,41 +123,39 @@ def generate_image(prompt, negative_prompt, steps, cfg, width, height):
     return image
 
 # -----------------------------------------------------------------------------
-# 3. INTERFAZ GRADIO
+# 3. INTERFAZ
 # -----------------------------------------------------------------------------
 css = """
 .container { max-width: 900px; margin: auto; }
 """
-
 DEFAULT_PROMPT = """<character_1>
 <gender>1girl</gender>
-<appearance>blue_eyes, silver_hair, long_hair</appearance>
-<clothing>school_uniform, serafuku</clothing>
-<action>standing, smile</action>
+<appearance>red_eyes, silver_hair, long_hair</appearance>
+<clothing>kimono, floral_print</clothing>
+<action>standing, holding_fan</action>
 </character_1>
 <general_tags>
-<quality>masterpiece, best quality, 4k</quality>
+<quality>best quality, masterpiece, 4k</quality>
 <style>anime, vivid_colors</style>
 </general_tags>"""
 
-DEFAULT_NEGATIVE = "low quality, bad anatomy, worst quality, watermark, text"
+DEFAULT_NEG = "low quality, bad anatomy, worst quality, watermark, text"
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    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.")
+    gr.Markdown("# ⛩️ NewBie Anime Generator (Hard-Import Fix)")
     
     with gr.Row():
         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)
+            prompt = gr.Textbox(label="Prompt (XML)", value=DEFAULT_PROMPT, lines=10)
+            neg = gr.Textbox(label="Negative", value=DEFAULT_NEG)
+            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")
-            btn = gr.Button("Generar", variant="primary")
         with gr.Column():
             out = gr.Image(label="Resultado")
 
-    btn.click(generate_image, inputs=[prompt_input, neg_prompt_input, steps, cfg], outputs=out)
+    btn.click(generate_image, inputs=[prompt, neg, steps, cfg], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()