Update app.py

app.py

@@ -26,48 +26,33 @@ pipe.to("cuda")
 pipe_upsample.to("cuda")
 pipe.vae.enable_tiling()
 
-def prepare_image_condition(image, size=(512, 512), background=(0, 0, 0)):
+def prepare_image_condition(image, size=(480, 480), background=(0, 0, 0)):
     image = ImageOps.contain(image, size)
     canvas = Image.new("RGB", size, background)
     offset = ((size[0] - image.width) // 2, (size[1] - image.height) // 2)
     canvas.paste(image, offset)
     return canvas
 
-def round_to_nearest_resolution(height, width, ratio):
-    return height - (height % ratio), width - (width % ratio)
-
 @spaces.GPU(duration=180)
 def generate_video(prompt, image_url):
     generator = torch.Generator("cuda").manual_seed(42)
-
-    # Aspect-ratio preserving image prep
+    
+    # Load & prepare image
     image = None
     if image_url:
         raw_image = Image.open(BytesIO(requests.get(image_url).content)).convert("RGB")
         image = prepare_image_condition(raw_image)
-
-    # Dimensions
-    base_width, base_height = 512, 512
-    downscale = 2 / 3
-    # Use correct rounding for VAE compatibility
-    w_d, h_d = round_to_nearest_resolution(
-        int(base_width * downscale),
-        int(base_height * downscale),
-        ratio=pipe.vae_spatial_compression_ratio
-    )
-    # Upscaled dimensions must also be VAE-aligned
-    w_up, h_up = round_to_nearest_resolution(
-        base_width,
-        base_height,
-        ratio=pipe.vae_spatial_compression_ratio
-    )
-
-    # Step 1: Generate latents
+    
+    # Set target resolutions
+    base_width, base_height = 480, 480  # final size (must be divisible by 16)
+    down_width, down_height = 320, 320  # for latent generation (must also be divisible by 16)
+    
+    # Step 1: Generate latents at lower resolution
     latents = pipe(
         prompt=prompt,
         image=image,
-        width=w_d,
-        height=h_d,
+        width=down_width,
+        height=down_height,
         num_frames=60,
         num_inference_steps=7,
         output_type="latent",
@@ -76,61 +61,98 @@ def generate_video(prompt, image_url):
         decode_noise_scale=0.025,
         generator=generator
     ).frames
-
+    
     torch.cuda.empty_cache()
     gc.collect()
-
-    # Step 2: Upscale
-    upscaled = pipe_upsample(latents=latents, output_type="latent").frames
-
+    
+    # Step 2: Upscale latents
+    upscaled_latents = pipe_upsample(latents=latents, output_type="latent").frames
+    
     torch.cuda.empty_cache()
     gc.collect()
-
-    # Step 3: Decode to frames (must match rounded base)
-    frames = pipe(
-        prompt=prompt,
-        image=image,
-        latents=upscaled,
-        width=w_up,
-        height=h_up,
-        num_frames=60,
-        num_inference_steps=10,
-        output_type="pil",
-        guidance_scale=1.0,
-        decode_timestep=0.05,
-        decode_noise_scale=0.025,
-        image_cond_noise_scale=0.025,
-        denoise_strength=0.3,
-        generator=generator
-    ).frames[0]
-
+    
+    # Step 3: Decode upscaled latents to frames
+    # Use the VAE decoder directly instead of the full pipeline
+    frames = pipe.vae.decode(upscaled_latents).sample
+    frames = (frames / 2 + 0.5).clamp(0, 1)  # Normalize to [0, 1]
+    frames = (frames * 255).to(torch.uint8)  # Convert to uint8
+    
+    # Convert tensor to PIL Images
+    pil_frames = []
+    for i in range(frames.shape[2]):  # num_frames dimension
+        frame = frames[0, :, i, :, :].permute(1, 2, 0).cpu().numpy()
+        pil_frames.append(Image.fromarray(frame))
+    
+    torch.cuda.empty_cache()
+    gc.collect()
+    
     # Step 4: Export video
     video_path = "output.mp4"
-    export_to_video(frames, video_path, fps=24)
-
+    export_to_video(pil_frames, video_path, fps=24)
+    
     # Step 5: TTS
     tts = gTTS(text=prompt, lang='en')
     tts.save("voice.mp3")
     AudioSegment.from_mp3("voice.mp3").export("voice.wav", format="wav")
-
-    # Step 6: Subtitles (CPU)
+    
+    # Step 6: Subtitles
     model = whisper.load_model("base", device="cpu")
     result = model.transcribe("voice.wav", task="transcribe", language="en")
+    
+    # Generate SRT subtitles manually since result["srt"] might not be available
+    srt_content = ""
+    for i, segment in enumerate(result["segments"]):
+        start_time = format_time(segment["start"])
+        end_time = format_time(segment["end"])
+        text = segment["text"].strip()
+        srt_content += f"{i + 1}\n{start_time} --> {end_time}\n{text}\n\n"
+    
     with open("subtitles.srt", "w", encoding="utf-8") as f:
-        f.write(result["srt"])
-
+        f.write(srt_content)
+    
     # Step 7: Merge video + audio + subtitles
     final_output = "final_with_audio.mp4"
-    ffmpeg.input(video_path).output(
-        final_output,
-        vf="subtitles=subtitles.srt",
-        i="voice.mp3",
-        c="copy",
-        shortest=None,
-        loglevel="error"
-    ).run()
-
-    return final_output
+    try:
+        (
+            ffmpeg
+            .input(video_path)
+            .output(
+                final_output,
+                vf="subtitles=subtitles.srt",
+                **{"c:v": "libx264", "c:a": "aac"},
+                loglevel="error"
+            )
+            .run(overwrite_output=True)
+        )
+        
+        # Add audio track
+        (
+            ffmpeg
+            .input(final_output)
+            .input("voice.wav")
+            .output(
+                "final_complete.mp4",
+                **{"c:v": "copy", "c:a": "aac"},
+                shortest=None,
+                loglevel="error"
+            )
+            .run(overwrite_output=True)
+        )
+        
+        return "final_complete.mp4"
+        
+    except Exception as e:
+        print(f"FFmpeg error: {e}")
+        # Fallback: return video without audio/subtitles
+        return video_path
+
+def format_time(seconds):
+    """Convert seconds to SRT time format"""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millisecs = int((seconds % 1) * 1000)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millisecs:03d}"
 
 # Gradio UI
 demo = gr.Interface(
@@ -141,7 +163,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Video(label="Generated Video"),
     title="🎬 LTX AI Video Generator",
-    description="AI-powered video with voiceover and subtitles. Supports ZeroGPU (PyTorch) runtime."
+    description="AI-powered video with voiceover and subtitles. Now outputs at 480x480 resolution."
 )
 
-demo.launch()
+demo.launch()