Update app.py

app.py

@@ -3,102 +3,73 @@ import torch
 import spaces
 import os
 import shutil
-import sys
-import importlib.util
+from diffusers import DiffusionPipeline
 from huggingface_hub import snapshot_download
 
 # -----------------------------------------------------------------------------
-# 1. FUNCIÓN DE INGENIERÍA INVERSA Y CARGA
+# 1. FUNCIÓN DE REPARACIÓN DE ESTRUCTURA (LA SOLUCIÓN)
 # -----------------------------------------------------------------------------
 MODEL_ID = "NewBie-AI/NewBie-image-Exp0.1"
 LOCAL_DIR = "./newbie_fixed_model"
 
-def load_hard_fixed_pipeline():
-    print(f"🛠️ Iniciando protocolo de carga manual para {MODEL_ID}...")
+def load_patched_pipeline():
+    print(f"🛠️ Iniciando protocolo de reparación v3 para {MODEL_ID}...")
     
-    # --- PASO 1: Descarga ---
+    # 1. Descargar repositorio completo
     if not os.path.exists(LOCAL_DIR):
-        print(f"   Descargando snapshot del modelo...")
-        snapshot_download(
-            repo_id=MODEL_ID,
-            local_dir=LOCAL_DIR,
-            ignore_patterns=["*.msgpack", "*.bin"]
-        )
+        print("   ⬇️ Descargando snapshot (esto puede tardar unos minutos)...")
+        snapshot_download(repo_id=MODEL_ID, local_dir=LOCAL_DIR)
     
-    # --- PASO 2: Arreglar estructura de carpetas (Transformer) ---
+    # 2. CREAR CARPETA TRANSFORMER (El error original)
     transformer_folder = os.path.join(LOCAL_DIR, "transformer")
-    if not os.path.exists(transformer_folder):
-        os.makedirs(transformer_folder, exist_ok=True)
-        # Buscar el archivo del transformer
-        candidates = [f for f in os.listdir(LOCAL_DIR) if f.endswith(".py") and "test" not in f]
-        for f in candidates:
-            # 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")
+    os.makedirs(transformer_folder, exist_ok=True)
     
-    # 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")]
+    # 3. BUSCAR Y MOVER ARCHIVOS PERDIDOS
+    # El config dice que el código está en 'transformer/transformer.py', pero suele estar en la raíz.
+    # Vamos a copiar TODOS los archivos .py de la raíz a la carpeta transformer por si acaso.
+    root_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()
-            
-            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
+    candidate_found = False
+    for f in root_files:
+        src = os.path.join(LOCAL_DIR, f)
+        
+        # Copia de seguridad: Copiamos todo .py a la carpeta transformer
+        shutil.copy(src, os.path.join(transformer_folder, f))
+        
+        # DETECTIVE: Leemos el archivo para ver si es el modelo principal
+        with open(src, "r", encoding="utf-8", errors="ignore") as file_content:
+            content = file_content.read()
+            # Si el archivo define un Transformer o un Modelo, es nuestro candidato
+            if "class" in content and ("Transformer" in content or "Model" in content):
+                # Lo copiamos como 'transformer.py' que es lo que pide el json
+                shutil.copy(src, os.path.join(transformer_folder, "transformer.py"))
+                print(f"   ✅ Candidato detectado y renombrado: {f} -> transformer/transformer.py")
+                candidate_found = True
 
-    if pipeline_class is None:
-        raise ValueError("❌ No se encontró la clase 'NewbiePipeline' en ningún archivo .py del repositorio.")
+    # Crear __init__.py obligatorio
+    with open(os.path.join(transformer_folder, "__init__.py"), "w") as f: f.write("")
 
-    # --- PASO 4: Carga del Modelo usando la Clase Importada ---
-    print("   🚀 Cargando pesos usando la clase personalizada...")
-    pipe = pipeline_class.from_pretrained(
+    if not candidate_found:
+        print("   ⚠️ No se detectó un modelo obvio. Copiando el archivo .py más grande como transformer.py")
+        # Fallback: Copiar el archivo .py más grande
+        largest_py = max(root_files, key=lambda x: os.path.getsize(os.path.join(LOCAL_DIR, x)))
+        shutil.copy(os.path.join(LOCAL_DIR, largest_py), os.path.join(transformer_folder, "transformer.py"))
+
+    # 4. CARGAR MODELO DESDE CARPETA LOCAL ARREGLADA
+    print("   🚀 Cargando pipeline...")
+    pipe = DiffusionPipeline.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 
+        trust_remote_code=True, # Necesario para ejecutar el código que acabamos de mover
+        local_files_only=True
     )
-    
     return pipe
 
-# Ejecutar carga
+# Carga inicial
 try:
-    pipe = load_hard_fixed_pipeline()
+    pipe = load_patched_pipeline()
 except Exception as e:
-    print(f"❌ ERROR CRÍTICO EN CARGA: {e}")
+    print(f"❌ Error fatal: {e}")
     pipe = None
 
 # -----------------------------------------------------------------------------
@@ -107,30 +78,37 @@ except Exception as e:
 @spaces.GPU(duration=120)
 def generate_image(prompt, negative_prompt, steps, cfg, width, height):
     if pipe is None:
-        raise gr.Error("El modelo no se pudo cargar. Revisa los logs.")
-        
+        raise gr.Error("El modelo no se cargó correctamente. Revisa los logs.")
+    
     print("🎨 Generando...")
     pipe.to("cuda")
     
-    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
+    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:
+        return None
 
 # -----------------------------------------------------------------------------
-# 3. INTERFAZ
+# 3. INTERFAZ (Corrección de CSS para Gradio 5.0)
 # -----------------------------------------------------------------------------
-css = """
+# Inyectamos CSS usando Markdown HTML en lugar del argumento 'css' que daba error
+custom_css = """
+<style>
 .container { max-width: 900px; margin: auto; }
+</style>
 """
+
 DEFAULT_PROMPT = """<character_1>
 <gender>1girl</gender>
-<appearance>red_eyes, silver_hair, long_hair</appearance>
+<appearance>red_eyes, white_hair, long_hair</appearance>
 <clothing>kimono, floral_print</clothing>
 <action>standing, holding_fan</action>
 </character_1>
@@ -141,8 +119,10 @@ DEFAULT_PROMPT = """<character_1>
 
 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 (Hard-Import Fix)")
+# Quitamos el argumento css=css aquí para evitar el TypeError
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.HTML(custom_css) # Solución al error de CSS
+    gr.Markdown("# ⛩️ NewBie Anime Generator (Fixed)")
     
     with gr.Row():
         with gr.Column():
@@ -152,10 +132,12 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
             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")
         with gr.Column():
             out = gr.Image(label="Resultado")
 
-    btn.click(generate_image, inputs=[prompt, neg, steps, cfg], outputs=out)
+    btn.click(generate_image, inputs=[prompt, neg, steps, cfg, width, height], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()