Update app_animatediff.py

app_animatediff.py

@@ -1,14 +1,14 @@
-# app_gradio_img2vid.py
 import os, io, tempfile
-from typing import Optional, List
+from typing import Optional
 from PIL import Image
 import torch
 import gradio as gr
 from diffusers import AnimateDiffPipeline, DDIMScheduler, MotionAdapter
 from diffusers.utils import export_to_gif
 
-MODEL_ID = "SG161222/Realistic_Vision_V5.1_noVAE"           # SD1.5 finetunado [attached_file:1]
-ADAPTER_ID = "guoyww/animatediff-motion-adapter-v1-5-2"     # MotionAdapter p/ SD1.4/1.5 [attached_file:1]
+# Modelos padrão (ajuste se desejar)
+MODEL_ID = "SG161222/Realistic_Vision_V5.1_noVAE"         # SD1.5 finetunado [attached_file:1]
+ADAPTER_ID = "guoyww/animatediff-motion-adapter-v1-5-2"   # MotionAdapter p/ SD1.4/1.5 [attached_file:1]
 
 pipe = None
 
@@ -16,19 +16,27 @@ def load_pipe(model_id: str, adapter_id: str, cpu_offload: bool):
     global pipe
     if pipe is not None:
         return pipe
+    # dtype preferível: float16 em CUDA, senão float32
     dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-    # 1) Carregar adapter SEM dtype
-    adapter = MotionAdapter.from_pretrained(adapter_id)  # antes: dtype=dtype (removido) [attached_file:1]
+    # MotionAdapter não aceita dtype em from_pretrained nas versões atuais
+    adapter = MotionAdapter.from_pretrained(adapter_id)  # [attached_file:1]
 
-    # 2) Carregar a pipeline com dtype
-    p = AnimateDiffPipeline.from_pretrained(
-        model_id,
-        motion_adapter=adapter,
-        dtype=dtype  # ou torch_dtype=dtype dependendo da versão instalada [attached_file:1]
-    )  # [attached_file:1]
+    # Carregar pipeline com dtype
+    try:
+        p = AnimateDiffPipeline.from_pretrained(
+            model_id,
+            motion_adapter=adapter,
+            dtype=dtype  # novas versões aceitam 'dtype' [attached_file:1]
+        )
+    except TypeError:
+        p = AnimateDiffPipeline.from_pretrained(
+            model_id,
+            motion_adapter=adapter,
+            torch_dtype=dtype  # fallback para versões que ainda usam torch_dtype [attached_file:1]
+        )
 
-    # 3) Scheduler recomendado
+    # Scheduler recomendado para estabilidade temporal
     p.scheduler = DDIMScheduler.from_pretrained(
         model_id,
         subfolder="scheduler",
@@ -38,23 +46,22 @@ def load_pipe(model_id: str, adapter_id: str, cpu_offload: bool):
         steps_offset=1
     )  # [attached_file:1]
 
-    # 4) Otimizações VAE nas APIs novas
-    p.vae.enable_slicing()   # [attached_file:1]
+    # Otimizações de VRAM (APIs novas via VAE)
+    p.vae.enable_slicing()  # [attached_file:1]
     try:
-        p.vae.enable_tiling()  # [attached_file:1]
+        p.vae.enable_tiling()  # útil em resoluções mais altas [attached_file:1]
     except Exception:
         pass
 
-    # 5) Device/offload
+    # Alocação de device / offload
     if cpu_offload and torch.cuda.is_available():
-        p.enable_model_cpu_offload()  # [attached_file:1]
+        p.enable_model_cpu_offload()  # reduz pico de VRAM [attached_file:1]
     else:
         p.to("cuda" if torch.cuda.is_available() else "cpu")
 
     pipe = p
     return pipe
 
-
 def generate(
     image: Image.Image,
     prompt: str,
@@ -72,9 +79,12 @@ def generate(
     cpu_offload: bool
 ):
     if image is None or not prompt or not prompt.strip():
-        return None, None, "Envie uma imagem e um prompt válidos."
+        return None, None, "Envie uma imagem e um prompt válidos."  # [attached_file:1]
+
     p = load_pipe(model_id_ui or MODEL_ID, adapter_id_ui or ADAPTER_ID, cpu_offload)
+
     gen = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))
+
     # img2vid sem IP-Adapter: NÃO passar ip_adapter_image
     out = p(
         prompt=prompt,
@@ -86,64 +96,61 @@ def generate(
         width=int(width) if width else None,
         height=int(height) if height else None
     )  # [attached_file:1]
-    frames = out.frames[0]  # lista de PIL [attached_file:1]
 
-    # GIF em memória
-    gif_buf = io.BytesIO()
-    export_to_gif(frames, gif_buf, fps=int(fps))  # [attached_file:1]
-    gif_buf.seek(0)
+    frames = out.frames[0]  # lista de PILs [attached_file:1]
+
+    # Salvar GIF em caminho temporário com extensão .gif (evita erro do PIL)
+    temp_gif = os.path.join(tempfile.gettempdir(), "animation.gif")
+    export_to_gif(frames, temp_gif, fps=int(fps))  # [attached_file:1]
 
+    # Opcional: gravar MP4 com imageio-ffmpeg
     mp4_path = None
     if save_mp4:
         try:
             import imageio
             mp4_path = os.path.join(tempfile.gettempdir(), "animation.mp4")
+            # Converter cada frame PIL para ndarray esperado pelo writer
             with imageio.get_writer(mp4_path, fps=int(fps), codec="libx264", quality=8) as writer:
                 for fr in frames:
-                    writer.append_data(imageio.v3.imread(io.BytesIO(fr.tobytes())))
+                    writer.append_data(imageio.v3.imread(io.BytesIO(fr.convert("RGB").tobytes())))
         except Exception:
-            mp4_path = None
+            mp4_path = None  # se falhar, apenas não retorna MP4
 
-    return gif_buf, mp4_path, f"Gerado {len(frames)} frames @ {fps} fps."
+    return temp_gif, mp4_path, f"Gerado {len(frames)} frames @ {fps} fps."  # [attached_file:1]
 
 def ui():
     with gr.Blocks(title="AnimateDiff img2vid") as demo:
-        gr.Markdown("## AnimateDiff img2vid")
+        gr.Markdown("## AnimateDiff img2vid")  # [attached_file:1]
         with gr.Row():
             with gr.Column(scale=1):
-                image = gr.Image(type="pil", label="Imagem inicial")
-                prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Descreva estilo/movimento...")
-                negative = gr.Textbox(label="Negative prompt", lines=2, value="low quality, worst quality")
+                image = gr.Image(type="pil", label="Imagem inicial")  # [attached_file:1]
+                prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Descreva estilo/movimento...")  # [attached_file:1]
+                negative = gr.Textbox(label="Negative prompt", lines=2, value="low quality, worst quality")  # [attached_file:1]
                 with gr.Row():
-                    frames = gr.Slider(8, 64, value=16, step=1, label="Frames")
-                    steps = gr.Slider(4, 60, value=25, step=1, label="Steps")
+                    frames = gr.Slider(8, 64, value=16, step=1, label="Frames")  # [attached_file:1]
+                    steps = gr.Slider(4, 60, value=25, step=1, label="Steps")  # [attached_file:1]
                 with gr.Row():
-                    guidance = gr.Slider(0.5, 15.0, value=7.5, step=0.5, label="Guidance")
-                    fps = gr.Slider(4, 30, value=8, step=1, label="FPS")
+                    guidance = gr.Slider(0.5, 15.0, value=7.5, step=0.5, label="Guidance")  # [attached_file:1]
+                    fps = gr.Slider(4, 30, value=8, step=1, label="FPS")  # [attached_file:1]
                 with gr.Row():
-                    seed = gr.Number(value=42, precision=0, label="Seed")
-                    width = gr.Number(value=None, precision=0, label="Largura (opcional)")
-                    height = gr.Number(value=None, precision=0, label="Altura (opcional)")
+                    seed = gr.Number(value=42, precision=0, label="Seed")  # [attached_file:1]
+                    width = gr.Number(value=None, precision=0, label="Largura (opcional)")  # [attached_file:1]
+                    height = gr.Number(value=None, precision=0, label="Altura (opcional)")  # [attached_file:1]
                 with gr.Row():
-                    model_id_ui = gr.Textbox(value=MODEL_ID, label="Model ID (SD1.5 finetune)")
-                    adapter_id_ui = gr.Textbox(value=ADAPTER_ID, label="MotionAdapter ID")
+                    model_id_ui = gr.Textbox(value=MODEL_ID, label="Model ID (SD1.5 finetune)")  # [attached_file:1]
+                    adapter_id_ui = gr.Textbox(value=ADAPTER_ID, label="MotionAdapter ID")  # [attached_file:1]
                 with gr.Row():
-                    cpu_offload = gr.Checkbox(value=False, label="CPU offload")
-                    save_mp4 = gr.Checkbox(value=False, label="Salvar MP4")
-                run_btn = gr.Button("Gerar animação")
+                    cpu_offload = gr.Checkbox(value=False, label="CPU offload")  # [attached_file:1]
+                    save_mp4 = gr.Checkbox(value=False, label="Salvar MP4")  # [attached_file:1]
+                run_btn = gr.Button("Gerar animação")  # [attached_file:1]
             with gr.Column(scale=1):
-                video_out = gr.Video(label="Preview (GIF salvo temporário)", interactive=False)
-                file_mp4 = gr.File(label="MP4 (download)", interactive=False)
-                status = gr.Textbox(label="Status", interactive=False)
+                video_out = gr.Video(label="Preview (GIF)")  # [attached_file:1]
+                file_mp4 = gr.File(label="MP4 (download)", interactive=False)  # [attached_file:1]
+                status = gr.Textbox(label="Status", interactive=False)  # [attached_file:1]
 
         def _run(*args):
-            gif_buf, mp4_path, msg = generate(*args)
-            temp_gif = None
-            if gif_buf:
-                temp_gif = os.path.join(tempfile.gettempdir(), "animation.gif")
-                with open(temp_gif, "wb") as f:
-                    f.write(gif_buf.read())
-            return temp_gif, mp4_path, msg
+            temp_gif, mp4_path, msg = generate(*args)
+            return temp_gif, mp4_path, msg  # [attached_file:1]
 
         run_btn.click(
             _run,
@@ -154,4 +161,4 @@ def ui():
 
 if __name__ == "__main__":
     demo = ui()
-    demo.launch(server_name="0.0.0.0", server_port=7860, inbrowser=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, inbrowser=True)  # [attached_file:1]