Update app.py

app.py

@@ -1,20 +1,22 @@
+```python
 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
+# Use 'balanced' for multi-device setups and CPU fallback for Spaces without GPU
+device_map = "balanced" if torch.cuda.is_available() else "cpu"
+
 flux_model = DiffusionPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.float16,
-    device_map="auto"
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map=device_map
 )
 
 omni_model = DiffusionPipeline.from_pretrained(
     "tencent/OmniAvatar",
-    torch_dtype=torch.float16,
-    device_map="auto"
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map=device_map
 )
 
 # ---------------------- MAIN GENERATION FUNCTION ----------------------
@@ -24,7 +26,7 @@ def generate_video(image, audio, prompt, style="claymation"):
     omni_model.to(device)
 
     try:
-        # Step 1: Stylize character image using FLUX-Kontext diffusion model
+        # Step 1: Stylize input image using FLUX-Kontext
         stylized_image = flux_model(
             prompt=prompt,
             image=image,
@@ -32,14 +34,14 @@ def generate_video(image, audio, prompt, style="claymation"):
             num_inference_steps=30
         ).images[0]
 
-        # Step 2: Animate the stylized image with lip-sync using OmniAvatar
+        # Step 2: Animate the stylized image with OmniAvatar
         result = omni_model(
             image=stylized_image,
             audio=audio,
             style=style,
         )
 
-        # The model should return a dictionary with "video" or similar key
+        # Return the generated video if available
         if isinstance(result, dict) and "video" in result:
             return result["video"]
         elif hasattr(result, "videos"):
@@ -85,3 +87,5 @@ with gr.Blocks(title="🎭 Claymation Talking Avatar Generator") as demo:
 
 # ---------------------- LAUNCH ----------------------
 demo.queue().launch(debug=True, share=False)
+```
+