Update app.py

app.py

@@ -6,112 +6,138 @@ import torch
 import numpy as np
 from PIL import Image
 
-# ── Load OVI Model ──────────────────────────────────────────
 print("Loading OVI model...")
 
-from transformers import AutoProcessor, AutoModel
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Device: {device}")
 
-MODEL_ID = "chetwinlow1/Ovi"
+model = None
+tokenizer = None
 
 try:
-    processor = AutoProcessor.from_pretrained(MODEL_ID)
-    model = AutoModel.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = model.to(device)
-    model.eval()
-    print(f"✅ OVI Model loaded on {device}")
+    from transformers import AutoTokenizer
+    from huggingface_hub import hf_hub_download, snapshot_download
+    import sys
+
+    # Download full repo
+    repo_path = snapshot_download("chetwinlow1/Ovi")
+    sys.path.insert(0, repo_path)
+
+    # Try importing model directly from repo
+    try:
+        from modeling_ovi import OviModel
+        from processing_ovi import OviProcessor
+
+        processor = OviProcessor.from_pretrained("chetwinlow1/Ovi")
+        model = OviModel.from_pretrained(
+            "chetwinlow1/Ovi",
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        ).to(device)
+        model.eval()
+        print("✅ OVI loaded via custom classes!")
+
+    except ImportError:
+        # Fallback - try pipeline
+        from transformers import pipeline
+        pipe = pipeline(
+            "image-to-video",
+            model="chetwinlow1/Ovi",
+            device=0 if device == "cuda" else -1,
+        )
+        model = pipe
+        processor = None
+        print("✅ OVI loaded via pipeline!")
+
 except Exception as e:
-    print(f"❌ Model load error: {e}")
-    processor = None
+    print(f"❌ Model error: {e}")
     model = None
+    processor = None
 
 
 def generate_video(image, prompt):
-    """Generate talking avatar video - FREE, no auth needed."""
-    
     if image is None:
         raise gr.Error("Please upload an image!")
-    
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter a text prompt!")
-    
     if model is None:
-        raise gr.Error("Model not loaded. Please restart the Space.")
-    
+        raise gr.Error("Model not loaded!")
+
     try:
-        # Process inputs
         if isinstance(image, str):
             pil_image = Image.open(image).convert("RGB")
         else:
             pil_image = Image.fromarray(image).convert("RGB")
-        
-        # Run model
-        inputs = processor(
-            text=prompt,
-            images=pil_image,
-            return_tensors="pt"
-        ).to(device)
-        
-        with torch.no_grad():
-            outputs = model.generate(**inputs)
-        
-        # Save output video
+
         output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-        
-        # Write video frames
+
+        if processor is not None:
+            # Custom processor path
+            inputs = processor(
+                text=prompt,
+                images=pil_image,
+                return_tensors="pt"
+            ).to(device)
+            with torch.no_grad():
+                outputs = model.generate(**inputs)
+        else:
+            # Pipeline path
+            outputs = model(pil_image, prompt)
+
+        # Save video
         if hasattr(outputs, 'video'):
-            video_data = outputs.video[0].cpu().numpy()
+            video_frames = outputs.video[0].cpu().numpy()
             import cv2
-            h, w = video_data.shape[1:3]
+            h, w = video_frames.shape[1:3]
             writer = cv2.VideoWriter(
                 output_path,
                 cv2.VideoWriter_fourcc(*'mp4v'),
                 24, (w, h)
             )
-            for frame in video_data:
+            for frame in video_frames:
                 frame_bgr = cv2.cvtColor(
                     (frame * 255).astype(np.uint8),
                     cv2.COLOR_RGB2BGR
                 )
                 writer.write(frame_bgr)
             writer.release()
-        
+        elif isinstance(outputs, str) and os.path.exists(outputs):
+            shutil.copy(outputs, output_path)
+        elif isinstance(outputs, list) and len(outputs) > 0:
+            out = outputs[0]
+            if isinstance(out, str) and os.path.exists(out):
+                shutil.copy(out, output_path)
+            elif hasattr(out, 'get'):
+                v = out.get('video') or out.get('path')
+                if v:
+                    shutil.copy(v, output_path)
+
         return output_path
-        
+
     except Exception as e:
-        raise gr.Error(f"Generation failed: {str(e)}")
+        raise gr.Error(f"Error: {str(e)}")
 
 
-# ── Gradio Interface ────────────────────────────────────────
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    
     gr.Markdown("""
 # 🎬 OVI — Talking Avatar Generator
-**Free & Open Source** | No login required | Powered by OVI AI Model
+**Free & Open Source** | No login required
 """)
-    
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
             image_input = gr.Image(
                 label="📸 Upload Image",
                 type="filepath",
-                sources=["upload", "clipboard"],
                 height=300,
             )
             prompt_input = gr.Textbox(
                 label="💬 Text Prompt",
                 lines=3,
-                placeholder=(
-                    "Describe what you want. Example:\n"
-                    "A person looks at camera and says "
-                    "<S>Hello everyone, welcome!<E> "
-                    "<AUDCAP>Clear friendly voice<ENDAUDCAP>"
-                ),
+                placeholder="A person speaks. <S>Hello world!<E> <AUDCAP>Clear voice<ENDAUDCAP>",
             )
             generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")
             clear_btn = gr.Button("🗑️ Clear", variant="secondary")
 
-        with gr.Column(scale=1):
+        with gr.Column():
             video_output = gr.Video(
                 label="🎥 Generated Video",
                 height=300,
@@ -119,30 +145,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
 
     gr.Markdown("""
----
-### 💡 Prompt Tips:
-- `<S>Your speech here<E>` — what avatar says
-- `<AUDCAP>voice description<ENDAUDCAP>` — voice style
-- Example: `A man speaks to camera. <S>Hello world!<E> <AUDCAP>Deep calm voice<ENDAUDCAP>`
+### 💡 Tips:
+- `<S>speech here<E>` — what avatar says
+- `<AUDCAP>voice style<ENDAUDCAP>` — voice description
 """)
 
-    gr.Examples(
-        examples=[
-            [None, "A person smiles at the camera. <S>Hello! Welcome to my channel.<E> <AUDCAP>Friendly energetic voice<ENDAUDCAP>"],
-            [None, "A woman speaks confidently. <S>Today I want to share something amazing with you.<E> <AUDCAP>Clear professional voice<ENDAUDCAP>"],
-        ],
-        inputs=[image_input, prompt_input],
-    )
-
     generate_btn.click(
         fn=generate_video,
         inputs=[image_input, prompt_input],
         outputs=[video_output],
     )
-    
     clear_btn.click(
         fn=lambda: (None, "", None),
-        inputs=None,
         outputs=[image_input, prompt_input, video_output],
     )