Create app.py

app.py

@@ -0,0 +1,87 @@
+import gradio as gr
+import torch
+from diffusers import DiffusionPipeline
+
+# ---------------------- MODEL INITIALIZATION ----------------------
+# Load Flux-Kontext and OmniAvatar pipelines from Hugging Face Hub
+# Both are large models — use torch_dtype and device_map for VRAM efficiency
+flux_model = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+
+omni_model = DiffusionPipeline.from_pretrained(
+    "tencent/OmniAvatar",
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+
+# ---------------------- MAIN GENERATION FUNCTION ----------------------
+def generate_video(image, audio, prompt, style="claymation"):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    flux_model.to(device)
+    omni_model.to(device)
+
+    try:
+        # Step 1: Stylize character image using FLUX-Kontext diffusion model
+        stylized_image = flux_model(
+            prompt=prompt,
+            image=image,
+            guidance_scale=7.5,
+            num_inference_steps=30
+        ).images[0]
+
+        # Step 2: Animate the stylized image with lip-sync using OmniAvatar
+        result = omni_model(
+            image=stylized_image,
+            audio=audio,
+            style=style,
+        )
+
+        # The model should return a dictionary with "video" or similar key
+        if isinstance(result, dict) and "video" in result:
+            return result["video"]
+        elif hasattr(result, "videos"):
+            return result.videos[0]
+        else:
+            return f"⚠️ Unexpected output format: {type(result)}"
+    except Exception as e:
+        return f"⚠️ Error during generation: {str(e)}"
+
+# ---------------------- GRADIO UI ----------------------
+with gr.Blocks(title="🎭 Claymation Talking Avatar Generator") as demo:
+    gr.Markdown("""
+    # 🎬 Claymation Talking Avatar Generator  
+    Generate claymation-style speaking avatars using **FLUX-Kontext** for stylization  
+    and **OmniAvatar** for lip-synced animation.
+    """)
+
+    with gr.Row():
+        image_input = gr.Image(label="🧑 Upload Character Image", type="filepath")
+        audio_input = gr.Audio(label="🎤 Upload Voice Audio", type="filepath")
+
+    prompt = gr.Textbox(
+        label="📝 Prompt (Optional)",
+        value="A claymation character speaking realistically",
+        placeholder="Describe the style or mood..."
+    )
+
+    with gr.Row():
+        style_dropdown = gr.Dropdown(
+            choices=["claymation", "toon", "realistic"],
+            value="claymation",
+            label="🎨 Style"
+        )
+        generate_button = gr.Button("🚀 Generate Video")
+
+    video_output = gr.Video(label="🎥 Generated Output")
+
+    generate_button.click(
+        fn=generate_video,
+        inputs=[image_input, audio_input, prompt, style_dropdown],
+        outputs=video_output
+    )
+
+# ---------------------- LAUNCH ----------------------
+demo.queue().launch(debug=True, share=False)