Update app.py

app.py

@@ -1,50 +1,61 @@
-import gradio as gr
 import torch
+import gradio as gr
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers import QwenImageEditPlusPipeline, FlowMatchEulerDiscreteScheduler
 import spaces
 
-# Usaremos un modelo que el sistema ya tiene en caché o que descarga de forma ligera
-MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
+# Configuración
+BASE_MODEL = "prithivMLmods/Qwen-Image-Edit-AIO-FP8"
+LORA_COSPLAY = "joyfox/Qwen-Image-Edit-Cosplay"
+LORA_LIGHTNING = "lightx2v/Qwen-Image-Lightning"
 
-# Esta configuración es mucho más amigable para el servidor gratuito
-pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-    MODEL_ID, 
-    torch_dtype=torch.float16, 
+# Inicialización básica del modelo
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16, # Usamos float16 para ahorrar VRAM
     use_safetensors=True
 )
 
-# Cargamos el inyector de estilo (IP-Adapter)
-pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.safetensors")
+# Cargamos los adaptadores de una vez al inicio
+pipe.load_lora_weights(LORA_COSPLAY, adapter_name="cosplay")
+pipe.load_lora_weights(LORA_LIGHTNING, weight_name="Qwen-Image-Lightning-8steps-V1.0.safetensors", adapter_name="lightning")
+pipe.set_adapters(["cosplay", "lightning"], adapter_weights=[0.85, 1.0])
 
 @spaces.GPU
-def generate_cosplay(base_image, reference_image, prompt):
-    # Movemos el pipeline a la GPU solo cuando se presiona el botón
+def process_cosplay_edit(base_image, reference_image, prompt, steps, cfg, seed):
+    # Movemos el modelo a la GPU activamente al iniciar la generación
     pipe.to("cuda")
-    pipe.set_ip_adapter_scale(0.7)
     
-    # Procesamiento
-    image = pipe(
+    # Redimensionamiento (múltiplos de 8)
+    width, height = 1024, 1024
+    base_resized = base_image.resize((width, height), Image.Resampling.LANCZOS)
+    ref_resized = reference_image.resize((width, height), Image.Resampling.LANCZOS)
+
+    generator = torch.Generator("cuda").manual_seed(int(seed))
+    
+    output = pipe(
+        image=[base_resized, ref_resized],
         prompt=prompt,
-        image=base_image.resize((1024, 1024)),
-        ip_adapter_image=reference_image.resize((1024, 1024)),
-        strength=0.7,
-        num_inference_steps=20
-    ).images[0]
+        width=width,
+        height=height,
+        num_inference_steps=int(steps),
+        true_cfg_scale=float(cfg),
+        generator=generator
+    )
     
-    # Liberamos la GPU inmediatamente después
-    pipe.to("cpu")
-    return image
+    # Al finalizar la función, el modelo se retira de la GPU automáticamente por ZeroGPU
+    return output.images[0]
 
-# Interfaz simplificada para evitar errores de memoria
+# Interfaz
 with gr.Blocks() as demo:
     with gr.Row():
-        base_img = gr.Image(label="Persona (Foto 1)", type="pil")
-        ref_img = gr.Image(label="Ropa (Foto 2)", type="pil")
-    prompt = gr.Textbox(value="masterpiece, realistic photo, person wearing the outfit from reference", label="Prompt")
-    btn = gr.Button("Generar Resultado")
-    result = gr.Image()
+        base_input = gr.Image(label="Persona Real", type="pil")
+        ref_input = gr.Image(label="Referencia Anime", type="pil")
     
-    btn.click(generate_cosplay, inputs=[base_img, ref_img, prompt], outputs=result)
+    prompt_input = gr.Textbox(label="Prompt", value="high quality, cosplay style transfer")
+    generate_btn = gr.Button("Generar", variant="primary")
+    output_display = gr.Image(label="Resultado")
+
+    generate_btn.click(process_cosplay_edit, inputs=[base_input, ref_input, prompt_input, gr.Slider(4, 12, 8), gr.Slider(1, 5, 1.5), gr.Number(1337)], outputs=output_display)
 
-demo.launch()
+demo.queue().launch()