Update app.py

app.py

@@ -2,15 +2,6 @@ import os
 import multiprocessing
 import subprocess
 import nltk
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from moviepy.editor import VideoFileClip
-import moviepy.editor as mpy
-from PIL import Image, ImageDraw, ImageFont
-from mutagen.mp3 import MP3
-from gtts import gTTS
-from pydub import AudioSegment
-import textwrap
 import gradio as gr
 import matplotlib.pyplot as plt
 import gc
@@ -19,31 +10,50 @@ from typing import List
 import shutil
 import numpy as np
 import random
+
+# Ensure `spaces` is imported first
+try:
+    import spaces
+except ImportError:
+    class spaces:
+        @staticmethod
+        def GPU(func=None, duration=None):
+            def wrapper(fn):
+                return fn
+            return wrapper if func is None else wrapper(func)
+
+# Now import CUDA-related libraries
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import DiffusionPipeline
-import spaces
-# Initialize FLUX pipeline outside of the function
+from moviepy.editor import VideoFileClip
+import moviepy.editor as mpy
+from PIL import Image, ImageDraw, ImageFont
+from mutagen.mp3 import MP3
+from gtts import gTTS
+from pydub import AudioSegment
+import textwrap
+
+# Initialize FLUX pipeline only if CUDA is available
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+
+if device == "cuda":
+    flux_pipe = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell", 
+        torch_dtype=dtype
+    ).to(device)
+else:
+    flux_pipe = None  # Avoid initializing the model when CUDA is unavailable
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 nltk.download('punkt')
+
 # Ensure proper multiprocessing start method
 multiprocessing.set_start_method("spawn", force=True)
 
-# GPU Fallback Setup
-#if os.environ.get("SPACES_ZERO_GPU") is not None:
-#    import spaces
-#else:
-#    class spaces:
-#        @staticmethod
-#        def GPU(func=None, duration=None):
-##            def wrapper(fn):
-#                return fn
-#            return wrapper if func is None else wrapper(func)
-
 # Download necessary NLTK data
 def setup_nltk():
     """Ensure required NLTK data is available."""
@@ -52,8 +62,6 @@ def setup_nltk():
     except LookupError:
         nltk.download('punkt')
 
-setup_nltk()
-
 # Constants
 DESCRIPTION = (
     "Video Story Generator with Audio\n"