Update app.py - full NSFW studio

app.py

@@ -1,75 +1,241 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
-from huggingface_hub import hf_hub_download
-
-# Cargamos el modelo base - usaremos SDXL como base
-# Podés cambiar esto por cualquier modelo de HuggingFace
-MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
-pipe = None
-
-def load_pipeline():
-    global pipe
-    if pipe is None:
-        pipe = DiffusionPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16"
+import numpy as np
+from PIL import Image
+import tempfile, os
+
+# ─── MODELOS ───────────────────────────────────────────────────────────────────
+TEXT2IMG_MODEL  = "SG161222/RealVisXL_V4.0"          # Realista, sin censura
+IMG2IMG_MODEL   = "SG161222/RealVisXL_V4.0"
+LTX_MODEL       = "Lightricks/LTX-Video"
+LTX_LORA        = "Lora-Daddy/Ltx2.3-real-nudity-early-alpha-30k-steps"
+
+pipe_t2i   = None
+pipe_i2i   = None
+pipe_video = None
+
+NEG_DEFAULT = (
+    "censored, blurry, low quality, bad anatomy, deformed, ugly, "
+    "watermark, logo, text, worst quality, jpeg artifacts"
+)
+
+# ─── LOADERS ───────────────────────────────────────────────────────────────────
+def load_t2i():
+    global pipe_t2i
+    if pipe_t2i is None:
+        from diffusers import StableDiffusionXLPipeline
+        pipe_t2i = StableDiffusionXLPipeline.from_pretrained(
+            TEXT2IMG_MODEL, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
+        )
+    return pipe_t2i
+
+def load_i2i():
+    global pipe_i2i
+    if pipe_i2i is None:
+        from diffusers import StableDiffusionXLImg2ImgPipeline
+        pipe_i2i = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            IMG2IMG_MODEL, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
+        )
+    return pipe_i2i
+
+def load_video(use_lora=True):
+    global pipe_video
+    if pipe_video is None:
+        from diffusers import LTXPipeline
+        pipe_video = LTXPipeline.from_pretrained(
+            LTX_MODEL, torch_dtype=torch.bfloat16
         )
-    return pipe
+        if use_lora:
+            try:
+                pipe_video.load_lora_weights(LTX_LORA)
+                print("[OK] LoRA NSFW cargado.")
+            except Exception as e:
+                print(f"[WARN] LoRA no cargado: {e}")
+    return pipe_video
+
+# ─── FUNCIONES ─────────────────────────────────────────────────────────────────
+@spaces.GPU(duration=120)
+def text2img(prompt, neg, steps, cfg, w, h, seed):
+    pipe = load_t2i().to("cuda")
+    gen = torch.Generator("cuda").manual_seed(int(seed))
+    img = pipe(prompt=prompt, negative_prompt=neg,
+                num_inference_steps=int(steps), guidance_scale=cfg,
+                width=int(w), height=int(h), generator=gen).images[0]
+    pipe.to("cpu"); torch.cuda.empty_cache()
+    return img
+
+@spaces.GPU(duration=120)
+def img2img(prompt, neg, init_image, strength, steps, cfg, seed):
+    if init_image is None:
+        return None
+    pipe = load_i2i().to("cuda")
+    gen = torch.Generator("cuda").manual_seed(int(seed))
+    img = Image.fromarray(init_image).convert("RGB").resize((1024, 1024))
+    result = pipe(prompt=prompt, negative_prompt=neg, image=img,
+                  strength=strength, num_inference_steps=int(steps),
+                  guidance_scale=cfg, generator=gen).images[0]
+    pipe.to("cpu"); torch.cuda.empty_cache()
+    return result
+
+@spaces.GPU(duration=200)
+def img2video(prompt, neg, init_image, num_frames, fps, steps, lora_scale, seed):
+    from diffusers.utils import export_to_video
+    pipe = load_video(use_lora=True).to("cuda")
+    gen = torch.Generator("cuda").manual_seed(int(seed))
+
+    kwargs = dict(
+        prompt=prompt, negative_prompt=neg,
+        num_frames=int(num_frames), num_inference_steps=int(steps),
+        generator=gen,
+    )
+    if hasattr(pipe, "image") and init_image is not None:
+        img = Image.fromarray(init_image).convert("RGB").resize((768, 512))
+        kwargs["image"] = img
+    if lora_scale > 0:
+        kwargs["cross_attention_kwargs"] = {"scale": lora_scale}
+
+    output = pipe(**kwargs)
+    frames = output.frames[0]
+
+    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    export_to_video(frames, tmp.name, fps=int(fps))
+    pipe.to("cpu"); torch.cuda.empty_cache()
+    return tmp.name
 
 @spaces.GPU(duration=120)
-def generate(prompt, negative_prompt, steps, cfg_scale, width, height, seed):
-    global pipe
-    pipe = load_pipeline()
-    pipe = pipe.to("cuda")
-    
-    generator = torch.Generator("cuda").manual_seed(int(seed))
-    
-    images = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        num_inference_steps=int(steps),
-        guidance_scale=cfg_scale,
-        width=int(width),
-        height=int(height),
-        generator=generator
-    ).images
-    
-    pipe = pipe.to("cpu")  # Liberar VRAM despues de generar
-    torch.cuda.empty_cache()
-    
-    return images[0]
-
-with gr.Blocks(title="Studio Privado", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🎨 Studio Privado - Generador de Imágenes")
-    gr.Markdown("*Tus creaciones son privadas. Nadie más puede verlas.*")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Prompt", placeholder="Describe lo que querés generar...", lines=3)
-            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, low quality, bad anatomy", lines=2)
-            
+def text2video(prompt, neg, num_frames, fps, w, h, steps, lora_scale, seed):
+    from diffusers import LTXPipeline
+    from diffusers.utils import export_to_video
+    pipe = load_video(use_lora=True).to("cuda")
+    gen = torch.Generator("cuda").manual_seed(int(seed))
+
+    kwargs = dict(
+        prompt=prompt, negative_prompt=neg,
+        num_frames=int(num_frames), width=int(w), height=int(h),
+        num_inference_steps=int(steps), generator=gen,
+    )
+    if lora_scale > 0:
+        kwargs["cross_attention_kwargs"] = {"scale": lora_scale}
+
+    output = pipe(**kwargs)
+    frames = output.frames[0]
+    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    export_to_video(frames, tmp.name, fps=int(fps))
+    pipe.to("cpu"); torch.cuda.empty_cache()
+    return tmp.name
+
+# ─── UI ────────────────────────────────────────────────────────────────────────
+THEME = gr.themes.Base(
+    primary_hue="violet", secondary_hue="purple", neutral_hue="slate"
+).set(
+    body_background_fill="#0f0f1a",
+    block_background_fill="#1a1a2e",
+    block_border_color="#7c3aed",
+    input_background_fill="#16213e",
+    button_primary_background_fill="linear-gradient(135deg, #7c3aed, #db2777)",
+    button_primary_text_color="white",
+)
+
+CSS = """
+h1 { text-align:center; color:#c084fc; font-size:2rem; margin-bottom:4px; }
+.subtitle { text-align:center; color:#94a3b8; margin-bottom:1rem; }
+.tab-nav button { font-weight:600; }
+"""
+
+with gr.Blocks(theme=THEME, css=CSS, title="Studio Privado NSFW") as demo:
+    gr.HTML("<h1>🔥 Studio Privado</h1>")
+    gr.HTML('<p class="subtitle">Generador multimedia sin censura · Tus creaciones son privadas y solo las ves vos</p>')
+
+    with gr.Tabs():
+        # ── TEXT → IMAGE ──────────────────────────────────────────────────────
+        with gr.Tab("🖼 Text → Image"):
             with gr.Row():
-                steps = gr.Slider(10, 50, value=30, step=1, label="Pasos")
-                cfg = gr.Slider(1, 20, value=7.5, step=0.5, label="CFG Scale")
-            
+                with gr.Column():
+                    t2i_prompt = gr.Textbox(label="Prompt", lines=3,
+                        placeholder="A beautiful woman, photorealistic, 8k, detailed...")
+                    t2i_neg    = gr.Textbox(label="Negative Prompt", value=NEG_DEFAULT, lines=2)
+                    with gr.Row():
+                        t2i_steps = gr.Slider(10, 60, 30, step=1, label="Pasos")
+                        t2i_cfg   = gr.Slider(1, 20, 7.5, step=0.5, label="CFG")
+                    with gr.Row():
+                        t2i_w   = gr.Slider(512, 1280, 1024, step=64, label="Ancho")
+                        t2i_h   = gr.Slider(512, 1280, 1024, step=64, label="Alto")
+                    t2i_seed  = gr.Number(42, label="Seed")
+                    t2i_btn   = gr.Button("🚀 Generar Imagen", variant="primary", size="lg")
+                with gr.Column():
+                    t2i_out = gr.Image(label="Resultado", type="pil", height=500)
+            t2i_btn.click(text2img,
+                [t2i_prompt, t2i_neg, t2i_steps, t2i_cfg, t2i_w, t2i_h, t2i_seed],
+                t2i_out)
+
+        # ── IMAGE → IMAGE ─────────────────────────────────────────────────────
+        with gr.Tab("🔄 Image → Image"):
             with gr.Row():
-                width = gr.Slider(512, 1024, value=1024, step=64, label="Ancho")
-                height = gr.Slider(512, 1024, value=1024, step=64, label="Alto")
-            
-            seed = gr.Number(value=42, label="Seed (-1 para random)")
-            btn = gr.Button("🚀 Generar", variant="primary", size="lg")
-        
-        with gr.Column(scale=1):
-            output = gr.Image(label="Resultado", type="pil")
-    
-    btn.click(
-        fn=generate,
-        inputs=[prompt, negative_prompt, steps, cfg, width, height, seed],
-        outputs=output
-    )
+                with gr.Column():
+                    i2i_input  = gr.Image(label="Imagen Base", type="numpy")
+                    i2i_prompt = gr.Textbox(label="Prompt", lines=3,
+                        placeholder="Modify the image to...")
+                    i2i_neg    = gr.Textbox(label="Negative Prompt", value=NEG_DEFAULT, lines=2)
+                    with gr.Row():
+                        i2i_str   = gr.Slider(0.1, 1.0, 0.6, step=0.05, label="Intensidad")
+                        i2i_steps = gr.Slider(10, 60, 30, step=1, label="Pasos")
+                    with gr.Row():
+                        i2i_cfg  = gr.Slider(1, 20, 7.5, step=0.5, label="CFG")
+                        i2i_seed = gr.Number(42, label="Seed")
+                    i2i_btn = gr.Button("🚀 Transformar Imagen", variant="primary", size="lg")
+                with gr.Column():
+                    i2i_out = gr.Image(label="Resultado", type="pil", height=500)
+            i2i_btn.click(img2img,
+                [i2i_prompt, i2i_neg, i2i_input, i2i_str, i2i_steps, i2i_cfg, i2i_seed],
+                i2i_out)
+
+        # ── TEXT → VIDEO (LTX) ────────────────────────────────────────────────
+        with gr.Tab("🎬 Text → Video (LTX)"):
+            gr.Markdown("*Usa LTX-Video con LoRA NSFW. Genera entre 25 y 121 frames.*")
+            with gr.Row():
+                with gr.Column():
+                    t2v_prompt = gr.Textbox(label="Prompt", lines=3,
+                        placeholder="A woman walking in slow motion, cinematic, 4k...")
+                    t2v_neg    = gr.Textbox(label="Negative Prompt", value=NEG_DEFAULT, lines=2)
+                    with gr.Row():
+                        t2v_frames = gr.Slider(25, 121, 49, step=8, label="Frames")
+                        t2v_fps    = gr.Slider(8, 30, 24, step=1, label="FPS")
+                    with gr.Row():
+                        t2v_w  = gr.Slider(256, 768, 512, step=64, label="Ancho")
+                        t2v_h  = gr.Slider(256, 768, 512, step=64, label="Alto")
+                    with gr.Row():
+                        t2v_steps = gr.Slider(10, 50, 30, step=1, label="Pasos")
+                        t2v_lora  = gr.Slider(0.0, 1.5, 0.8, step=0.05, label="LoRA Scale")
+                    t2v_seed = gr.Number(42, label="Seed")
+                    t2v_btn  = gr.Button("🎬 Generar Video", variant="primary", size="lg")
+                with gr.Column():
+                    t2v_out = gr.Video(label="Video Generado")
+            t2v_btn.click(text2video,
+                [t2v_prompt, t2v_neg, t2v_frames, t2v_fps, t2v_w, t2v_h, t2v_steps, t2v_lora, t2v_seed],
+                t2v_out)
+
+        # ── IMAGE → VIDEO (LTX) ───────────────────────────────────────────────
+        with gr.Tab("📹 Image → Video (LTX)"):
+            gr.Markdown("*Animá una imagen usando LTX-Video + LoRA NSFW.*")
+            with gr.Row():
+                with gr.Column():
+                    i2v_input  = gr.Image(label="Imagen Base (se usará como frame inicial)", type="numpy")
+                    i2v_prompt = gr.Textbox(label="Prompt de movimiento", lines=3,
+                        placeholder="The woman slowly turns her head, smooth motion...")
+                    i2v_neg    = gr.Textbox(label="Negative Prompt", value=NEG_DEFAULT, lines=2)
+                    with gr.Row():
+                        i2v_frames = gr.Slider(25, 121, 49, step=8, label="Frames")
+                        i2v_fps    = gr.Slider(8, 30, 24, step=1, label="FPS")
+                    with gr.Row():
+                        i2v_steps = gr.Slider(10, 50, 30, step=1, label="Pasos")
+                        i2v_lora  = gr.Slider(0.0, 1.5, 0.8, step=0.05, label="LoRA Scale")
+                    i2v_seed = gr.Number(42, label="Seed")
+                    i2v_btn  = gr.Button("📹 Animar Imagen", variant="primary", size="lg")
+                with gr.Column():
+                    i2v_out = gr.Video(label="Video Generado")
+            i2v_btn.click(img2video,
+                [i2v_prompt, i2v_neg, i2v_input, i2v_frames, i2v_fps, i2v_steps, i2v_lora, i2v_seed],
+                i2v_out)
 
 demo.launch()