Update app.py

app.py

@@ -17,13 +17,27 @@ from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from diffusers import (
     FluxPipeline,
-    FluxControlNetPipeline, 
     DDIMScheduler,
     DPMSolverMultistepScheduler
 )
 import soundfile as sf
 import requests
 
+# Optional imports for enhanced performance
+try:
+    import flash_attn
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+    print("⚠️ Flash Attention not available - using standard attention")
+
+try:
+    import triton
+    TRITON_AVAILABLE = True
+except ImportError:
+    TRITON_AVAILABLE = False
+    print("⚠️ Triton not available - using standard operations")
+
 class ProfessionalCartoonFilmGenerator:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -76,11 +90,43 @@ class ProfessionalCartoonFilmGenerator:
             self.flux_pipe.enable_vae_slicing()
             self.flux_pipe.enable_vae_tiling()
             
+            # Enable flash attention if available
+            if FLASH_ATTN_AVAILABLE:
+                try:
+                    self.flux_pipe.enable_xformers_memory_efficient_attention()
+                    print("✅ Flash attention enabled for better performance")
+                except Exception as e:
+                    print(f"⚠️ Flash attention failed: {e}")
+            else:
+                print("ℹ️ Using standard attention (flash attention not available)")
+            
             print("✅ FLUX pipeline loaded successfully")
             
         except Exception as e:
             print(f"❌ FLUX pipeline failed: {e}")
-            self.flux_pipe = None
+            print("🔄 Falling back to Stable Diffusion...")
+            
+            # Fallback to Stable Diffusion
+            try:
+                from diffusers import StableDiffusionPipeline
+                self.flux_pipe = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5",
+                    torch_dtype=torch.float16,
+                    use_safetensors=True,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                ).to(self.device)
+                
+                # Enable memory optimizations
+                self.flux_pipe.enable_vae_slicing()
+                if hasattr(self.flux_pipe, 'enable_vae_tiling'):
+                    self.flux_pipe.enable_vae_tiling()
+                
+                print("✅ Stable Diffusion fallback loaded successfully")
+                
+            except Exception as e2:
+                print(f"❌ Stable Diffusion fallback also failed: {e2}")
+                self.flux_pipe = None
         
         try:
             # 2. Advanced script generation model
@@ -356,17 +402,20 @@ class ProfessionalCartoonFilmGenerator:
         character_images = {}
         
         if not self.flux_pipe:
-            print("❌ FLUX pipeline not available")
+            print("❌ No image generation pipeline available")
             return character_images
         
         for character in characters:
             try:
                 print(f"🎭 Generating professional character: {character['name']}")
                 
-                # Load appropriate LoRA based on character type
-                if "anime" in character.get("animation_style", "").lower():
+                # Load appropriate LoRA based on character type (only for FLUX)
+                if hasattr(self.flux_pipe, 'load_lora_weights') and "anime" in character.get("animation_style", "").lower():
                     if hasattr(self, 'cartoon_lora'):
-                        self.flux_pipe.load_lora_weights(self.cartoon_lora)
+                        try:
+                            self.flux_pipe.load_lora_weights(self.cartoon_lora)
+                        except Exception as e:
+                            print(f"⚠️ LoRA loading failed: {e}")
                 
                 # Professional character prompt
                 prompt = f"""
@@ -382,15 +431,28 @@ class ProfessionalCartoonFilmGenerator:
                 inconsistent, amateur, simple, crude, manga, sketch
                 """
                 
-                image = self.flux_pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    num_inference_steps=25,  # High quality steps
-                    guidance_scale=3.5,
-                    height=1024,  # High resolution
-                    width=1024,
-                    max_sequence_length=256
-                ).images[0]
+                # Handle different pipeline types
+                if hasattr(self.flux_pipe, 'max_sequence_length'):
+                    # FLUX pipeline
+                    image = self.flux_pipe(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=25,  # High quality steps
+                        guidance_scale=3.5,
+                        height=1024,  # High resolution
+                        width=1024,
+                        max_sequence_length=256
+                    ).images[0]
+                else:
+                    # Stable Diffusion pipeline
+                    image = self.flux_pipe(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=25,  # High quality steps
+                        guidance_scale=7.5,
+                        height=1024,  # High resolution
+                        width=1024
+                    ).images[0]
                 
                 char_path = f"{self.temp_dir}/character_{character['name'].replace(' ', '_')}.png"
                 image.save(char_path)
@@ -411,7 +473,7 @@ class ProfessionalCartoonFilmGenerator:
         background_images = {}
         
         if not self.flux_pipe:
-            print("❌ FLUX pipeline not available")
+            print("❌ No image generation pipeline available")
             return background_images
         
         for scene in scenes:
@@ -432,15 +494,28 @@ class ProfessionalCartoonFilmGenerator:
                 blurry, simple, amateur, 3D render
                 """
                 
-                image = self.flux_pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    num_inference_steps=20,
-                    guidance_scale=3.0,
-                    height=768,   # 4:3 aspect ratio for traditional animation
-                    width=1024,
-                    max_sequence_length=256
-                ).images[0]
+                # Handle different pipeline types for backgrounds
+                if hasattr(self.flux_pipe, 'max_sequence_length'):
+                    # FLUX pipeline
+                    image = self.flux_pipe(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=20,
+                        guidance_scale=3.0,
+                        height=768,   # 4:3 aspect ratio for traditional animation
+                        width=1024,
+                        max_sequence_length=256
+                    ).images[0]
+                else:
+                    # Stable Diffusion pipeline
+                    image = self.flux_pipe(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=20,
+                        guidance_scale=7.0,
+                        height=768,   # 4:3 aspect ratio for traditional animation
+                        width=1024
+                    ).images[0]
                 
                 bg_path = f"{self.temp_dir}/background_scene_{scene['scene_number']}.png"
                 image.save(bg_path)
@@ -459,21 +534,36 @@ class ProfessionalCartoonFilmGenerator:
         try:
             print("🎬 Setting up Open-Sora 2.0 for video generation...")
             
-            # Clone Open-Sora repository
-            if not os.path.exists("Open-Sora"):
+            # Check if we're already in the right directory
+            current_dir = os.getcwd()
+            opensora_dir = os.path.join(current_dir, "Open-Sora")
+            
+            # Clone Open-Sora repository if it doesn't exist
+            if not os.path.exists(opensora_dir):
+                print("📥 Cloning Open-Sora repository...")
                 subprocess.run([
                     "git", "clone", "https://github.com/hpcaitech/Open-Sora.git"
                 ], check=True, capture_output=True)
             
-            os.chdir("Open-Sora")
-            
-            # Download model weights
-            print("📥 Downloading Open-Sora 2.0 model...")
-            subprocess.run([
-                "huggingface-cli", "download", "hpcai-tech/Open-Sora-v2", 
-                "--local-dir", "./ckpts"
-            ], check=True, capture_output=True)
-            
+            # Check if the repository was cloned successfully
+            if not os.path.exists(opensora_dir):
+                print("❌ Failed to clone Open-Sora repository")
+                return False
+            
+            # Check if model weights exist
+            ckpts_dir = os.path.join(opensora_dir, "ckpts")
+            if not os.path.exists(ckpts_dir):
+                print("📥 Downloading Open-Sora 2.0 model...")
+                try:
+                    subprocess.run([
+                        "huggingface-cli", "download", "hpcai-tech/Open-Sora-v2", 
+                        "--local-dir", ckpts_dir
+                    ], check=True, capture_output=True)
+                except Exception as e:
+                    print(f"❌ Model download failed: {e}")
+                    return False
+            
+            print("✅ Open-Sora setup completed")
             return True
             
         except Exception as e:
@@ -526,6 +616,14 @@ class ProfessionalCartoonFilmGenerator:
             
             video_path = f"{self.temp_dir}/scene_{scene['scene_number']}.mp4"
             
+            # Get the correct Open-Sora directory
+            current_dir = os.getcwd()
+            opensora_dir = os.path.join(current_dir, "Open-Sora")
+            
+            if not os.path.exists(opensora_dir):
+                print("❌ Open-Sora directory not found")
+                return None
+            
             # Run Open-Sora inference
             cmd = [
                 "torchrun", "--nproc_per_node", "1", "--standalone",
@@ -538,7 +636,7 @@ class ProfessionalCartoonFilmGenerator:
                 "--motion-score", "6"  # High motion for dynamic scenes
             ]
             
-            result = subprocess.run(cmd, capture_output=True, text=True, cwd="Open-Sora")
+            result = subprocess.run(cmd, capture_output=True, text=True, cwd=opensora_dir)
             
             if result.returncode == 0:
                 # Find generated video file