Update app.py

app.py

@@ -14,143 +14,92 @@ import gc
 from huggingface_hub import hf_hub_download
 import threading
 
-# Latest and best open-source models
+# ZeroGPU-compatible imports
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from diffusers import (
-    FluxPipeline,
+    StableDiffusionPipeline,
     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")
+# ZeroGPU compatibility - disable GPU-specific optimizations
+FLASH_ATTN_AVAILABLE = False
+TRITON_AVAILABLE = False
+print("⚠️ ZeroGPU mode - using CPU-optimized operations")
 
 # Global lock to prevent concurrent generations
 generation_lock = threading.Lock()
 
 class ProfessionalCartoonFilmGenerator:
     def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # ZeroGPU compatibility - force CPU usage
+        self.device = "cpu"
+        self.dtype = torch.float32  # Use float32 for CPU compatibility
         
         # Use /tmp directory for Hugging Face Spaces storage
-        # Note: HF Spaces only allows storage in /tmp, no subdirectories allowed
         self.output_dir = "/tmp"
         print(f"📁 Using Hugging Face temp directory: {self.output_dir}")
         
-        # Note: Cannot create subdirectories in /tmp on HF Spaces
-        # Will use file prefixes instead (char_, bg_, video_, etc.)
-        
         # Model configurations for ZeroGPU optimization
         self.models_loaded = False
-        self.using_flux = False
+        self.flux_available = False
         self.flux_pipe = None
-        self.script_enhancer = None
-        self.cartoon_lora = None
-        self.character_lora = None
-        self.sketch_lora = None
+        self.sd_pipe = None
+        self.script_model = None
+        self.script_tokenizer = None
         
-    @spaces.GPU
     def load_models(self):
-        """Load all required AI models for professional generation"""
+        """Load ZeroGPU-compatible models for professional generation"""
         try:
-            print("🚀 Loading professional-grade models...")
-            
-            # Clear GPU memory first
-            self.clear_gpu_memory()
+            print("🚀 Loading ZeroGPU-compatible models...")
             
-            # Detect device and set appropriate dtype
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-            self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+            # Clear memory
+            gc.collect()
             
             print(f"🎮 Using device: {self.device} with dtype: {self.dtype}")
             
-            # Try to load FLUX first
-            try:
-                print("🎨 Loading FLUX pipeline...")
-                from diffusers import FluxPipeline
-                
-                self.flux_pipe = FluxPipeline.from_pretrained(
-                    "black-forest-labs/FLUX.1-dev",
-                    torch_dtype=self.dtype,
-                    device_map="auto" if self.device == "cuda" else None
-                )
-                
-                if self.device == "cuda":
-                    self.flux_pipe = self.flux_pipe.to("cuda")
-                
-                print("✅ FLUX pipeline loaded successfully")
-                self.flux_available = True
-                
-            except Exception as e:
-                print("🔐 FLUX authentication failed - model requires Hugging Face token")
-                print("💡 To use FLUX, you need to:")
-                print("   1. Get a Hugging Face token from https://huggingface.co/settings/tokens")
-                print("   2. Accept the FLUX model license at https://huggingface.co/black-forest-labs/FLUX.1-dev")
-                print("   3. Set your token: huggingface-cli login")
-                print("🔄 Falling back to Stable Diffusion...")
-                self.flux_available = False
-            
-            # Load Stable Diffusion fallback
-            if not self.flux_available:
-                print("🔄 Loading Stable Diffusion fallback model...")
-                from diffusers import StableDiffusionPipeline, DDIMScheduler
-                
-                self.sd_pipe = StableDiffusionPipeline.from_pretrained(
-                    "CompVis/stable-diffusion-v1-4",
-                    torch_dtype=self.dtype,
-                    safety_checker=None,
-                    requires_safety_checker=False
-                )
-                
-                # Configure scheduler for better quality
-                self.sd_pipe.scheduler = DDIMScheduler.from_config(self.sd_pipe.scheduler.config)
-                
-                if self.device == "cuda":
-                    self.sd_pipe = self.sd_pipe.to("cuda")
+            # Load Stable Diffusion (CPU optimized)
+            print("🔄 Loading Stable Diffusion (CPU optimized)...")
+            from diffusers import StableDiffusionPipeline, DDIMScheduler
+            
+            self.sd_pipe = StableDiffusionPipeline.from_pretrained(
+                "CompVis/stable-diffusion-v1-4",
+                torch_dtype=self.dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
+                device_map=None  # Force CPU usage
+            )
+            
+            # Configure scheduler for better quality
+            self.sd_pipe.scheduler = DDIMScheduler.from_config(self.sd_pipe.scheduler.config)
+            
+            # Force CPU usage for ZeroGPU
+            self.sd_pipe = self.sd_pipe.to("cpu")
+            self.sd_pipe.enable_sequential_cpu_offload()  # Memory optimization
                     
-                print("✅ Loaded Stable Diffusion v1.4")
-                print("✅ Stable Diffusion fallback loaded successfully")
+            print("✅ Loaded Stable Diffusion v1.4 (CPU optimized)")
             
-            # Load script enhancement model with correct device
+            # Load script enhancement model (CPU optimized)
             print("📝 Loading script enhancement model...")
             self.script_model = AutoModelForCausalLM.from_pretrained(
                 "microsoft/DialoGPT-medium",
                 torch_dtype=self.dtype,
-                device_map="auto" if self.device == "cuda" else None
+                device_map=None  # Force CPU usage
             )
             self.script_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
             
             if self.script_tokenizer.pad_token is None:
                 self.script_tokenizer.pad_token = self.script_tokenizer.eos_token
             
-            if self.device == "cuda":
-                self.script_model = self.script_model.to("cuda")
+            # Force CPU usage
+            self.script_model = self.script_model.to("cpu")
             
             print(f"Device set to use {self.device}")
-            print("✅ Script enhancer loaded")
-            
-            # Set model states
-            if self.device == "cuda":
-                if self.flux_available:
-                    self.flux_pipe.enable_model_cpu_offload()
-                else:
-                    self.sd_pipe.enable_model_cpu_offload()
+            print("✅ Script enhancer loaded (CPU optimized)")
                     
-            print("🎬 All professional models loaded!")
+            print("🎬 All ZeroGPU-compatible models loaded!")
             return True
             
         except Exception as e:
@@ -160,10 +109,8 @@ class ProfessionalCartoonFilmGenerator:
             return False
     
     def clear_gpu_memory(self):
-        """Clear GPU memory between operations"""
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
+        """Clear memory (CPU-focused for ZeroGPU)"""
+        gc.collect()
     
     def optimize_prompt_for_clip(self, prompt: str, max_tokens: int = 70) -> str:
         """Optimize prompt to fit within CLIP token limit"""
@@ -201,8 +148,7 @@ class ProfessionalCartoonFilmGenerator:
             return download_info
             
         except Exception as e:
-            print(f"⚠️ Failed to create download info: {e}")
-            return f"📁 File generated: {file_path}"
+            return f"📥 Generated {file_type} (download info unavailable: {e})"
     
     def generate_professional_script(self, user_input: str) -> Dict[str, Any]:
         """Generate a professional cartoon script with detailed character development"""
@@ -447,28 +393,19 @@ class ProfessionalCartoonFilmGenerator:
         }
         return palettes.get(mood, "balanced warm and cool tones")
     
-    @spaces.GPU
     def generate_professional_character_images(self, characters: List[Dict]) -> Dict[str, str]:
-        """Generate professional character images with consistency"""
+        """Generate professional character images with consistency (ZeroGPU compatible)"""
         character_images = {}
         
         print(f"🎭 Generating {len(characters)} professional character designs...")
         
-        # Check if we have any image generation pipeline available
-        if not hasattr(self, 'flux_available'):
-            print("❌ No image generation models loaded")
+        # Check if we have Stable Diffusion pipeline available
+        if not hasattr(self, 'sd_pipe') or self.sd_pipe is None:
+            print("❌ Stable Diffusion not loaded - please call load_models() first")
             return character_images
             
-        pipeline = None
-        if self.flux_available and hasattr(self, 'flux_pipe'):
-            pipeline = self.flux_pipe
-            model_name = "FLUX"
-        elif hasattr(self, 'sd_pipe'):
-            pipeline = self.sd_pipe
-            model_name = "Stable Diffusion"
-        else:
-            print("❌ No image generation pipeline available")
-            return character_images
+        pipeline = self.sd_pipe
+        model_name = "Stable Diffusion (CPU)"
             
         print(f"🎨 Using {model_name} for character generation")
         
@@ -477,44 +414,28 @@ class ProfessionalCartoonFilmGenerator:
             print(f"\n🎨 Generating character: {character_name}")
             
             try:
-                # Build comprehensive character prompt
+                # Build comprehensive character prompt for CPU generation
                 base_prompt = f"Professional cartoon character design, {character['name']}, {character['description']}"
                 
-                # Add style and quality modifiers
-                if self.flux_available:
-                    # FLUX-specific prompt
-                    prompt = f"{base_prompt}, Disney-Pixar animation style, highly detailed character sheet, clean white background, 2D animation model sheet, expressive face, vibrant colors, professional character design, perfect for animation"
-                else:
-                    # Stable Diffusion prompt
-                    prompt = f"{base_prompt}, anime style, cartoon character, clean background, high quality, detailed, 2D animation style, character sheet"
+                # CPU-optimized prompt
+                prompt = f"{base_prompt}, anime style, cartoon character, clean background, high quality, detailed, 2D animation style, character sheet, simple design"
                 
                 # Optimize prompt for CLIP
-                prompt = self.optimize_prompt_for_clip(prompt, max_tokens=75)
+                prompt = self.optimize_prompt_for_clip(prompt, max_tokens=60)  # Shorter for CPU
                 print(f"📝 Character prompt: {prompt}")
                 
-                # Generate with appropriate settings
-                if self.flux_available:
-                    # FLUX generation settings
-                    image = pipeline(
-                        prompt=prompt,
-                        width=1024,
-                        height=1024,
-                        num_inference_steps=25,
-                        guidance_scale=7.5,
-                        generator=torch.Generator(device=self.device).manual_seed(42)
-                    ).images[0]
-                else:
-                    # Stable Diffusion generation settings
-                    image = pipeline(
-                        prompt=prompt,
-                        width=512,
-                        height=512,
-                        num_inference_steps=30,
-                        guidance_scale=7.5,
-                        generator=torch.Generator(device=self.device).manual_seed(42)
-                    ).images[0]
-                    # Upscale for SD
-                    image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
+                # CPU-optimized generation settings
+                image = pipeline(
+                    prompt=prompt,
+                    width=512,  # Smaller for CPU
+                    height=512,
+                    num_inference_steps=20,  # Fewer steps for CPU
+                    guidance_scale=7.5,
+                    generator=torch.Generator(device="cpu").manual_seed(42)
+                ).images[0]
+                
+                # Upscale for better quality
+                image = image.resize((1024, 1024), Image.Resampling.LANCZOS)
                 
                 # Save character image
                 char_path = f"{self.output_dir}/char_{character['name'].replace(' ', '_')}.png"
@@ -532,10 +453,8 @@ class ProfessionalCartoonFilmGenerator:
                     print(f"   📁 Internal path: {char_path}")
                     print(download_info)
                     
-                    # Clear GPU memory after each generation
-                    if self.device == "cuda":
-                        torch.cuda.empty_cache()
-                        gc.collect()
+                    # Clear memory after each generation
+                    gc.collect()
                 else:
                     print(f"❌ Failed to save character image: {char_path}")
                 
@@ -553,28 +472,19 @@ class ProfessionalCartoonFilmGenerator:
         
         return character_images
     
-    @spaces.GPU  
     def generate_cinematic_backgrounds(self, scenes: List[Dict], color_palette: str) -> Dict[int, str]:
-        """Generate professional cinematic backgrounds for each scene"""
+        """Generate professional cinematic backgrounds for each scene (ZeroGPU compatible)"""
         background_images = {}
         
         print(f"🎞️ Generating {len(scenes)} cinematic backgrounds...")
         
-        # Check if we have any image generation pipeline available
-        if not hasattr(self, 'flux_available'):
-            print("❌ No image generation models loaded")
+        # Check if we have Stable Diffusion pipeline available
+        if not hasattr(self, 'sd_pipe') or self.sd_pipe is None:
+            print("❌ Stable Diffusion not loaded - please call load_models() first")
             return background_images
             
-        pipeline = None
-        if self.flux_available and hasattr(self, 'flux_pipe'):
-            pipeline = self.flux_pipe
-            model_name = "FLUX"
-        elif hasattr(self, 'sd_pipe'):
-            pipeline = self.sd_pipe
-            model_name = "Stable Diffusion"
-        else:
-            print("❌ No image generation pipeline available")
-            return background_images
+        pipeline = self.sd_pipe
+        model_name = "Stable Diffusion (CPU)"
             
         print(f"🎨 Using {model_name} for background generation")
         
@@ -583,7 +493,7 @@ class ProfessionalCartoonFilmGenerator:
             print(f"\n🌄 Generating background for scene {scene_num}")
             
             try:
-                # Build cinematic background prompt
+                # Build cinematic background prompt for CPU generation
                 background_desc = scene['background']
                 mood = scene.get('mood', 'neutral')
                 shot_type = scene.get('shot_type', 'medium shot')
@@ -591,39 +501,25 @@ class ProfessionalCartoonFilmGenerator:
                 
                 base_prompt = f"Cinematic background scene, {background_desc}, {mood} atmosphere, {lighting}"
                 
-                # Add style and quality modifiers
-                if self.flux_available:
-                    prompt = f"{base_prompt}, Disney-Pixar animation style, detailed landscape, professional background art, vibrant colors, high quality, cinematic composition, no characters"
-                else:
-                    prompt = f"{base_prompt}, anime style background, detailed landscape, high quality, cinematic, {color_palette} color palette, no people"
+                # CPU-optimized prompt
+                prompt = f"{base_prompt}, anime style background, detailed landscape, high quality, cinematic, {color_palette} color palette, no people, simple design"
                 
                 # Optimize for CLIP
-                prompt = self.optimize_prompt_for_clip(prompt, max_tokens=75)
+                prompt = self.optimize_prompt_for_clip(prompt, max_tokens=60)  # Shorter for CPU
                 print(f"📝 Background prompt: {prompt}")
                 
-                # Generate with appropriate settings
-                if self.flux_available:
-                    # FLUX generation settings
-                    image = pipeline(
-                        prompt=prompt,
-                        width=1024,
-                        height=768,  # 4:3 aspect ratio for video
-                        num_inference_steps=25,
-                        guidance_scale=7.5,
-                        generator=torch.Generator(device=self.device).manual_seed(scene_num * 10)
-                    ).images[0]
-                else:
-                    # Stable Diffusion generation settings
-                    image = pipeline(
-                        prompt=prompt,
-                        width=512,
-                        height=384,  # 4:3 aspect ratio
-                        num_inference_steps=30,
-                        guidance_scale=7.5,
-                        generator=torch.Generator(device=self.device).manual_seed(scene_num * 10)
-                    ).images[0]
-                    # Upscale for SD
-                    image = image.resize((1024, 768), Image.Resampling.LANCZOS)
+                # CPU-optimized generation settings
+                image = pipeline(
+                    prompt=prompt,
+                    width=512,  # Smaller for CPU
+                    height=384,  # 4:3 aspect ratio
+                    num_inference_steps=20,  # Fewer steps for CPU
+                    guidance_scale=7.5,
+                    generator=torch.Generator(device="cpu").manual_seed(scene_num * 10)
+                ).images[0]
+                
+                # Upscale for better quality
+                image = image.resize((1024, 768), Image.Resampling.LANCZOS)
                 
                 # Save background image
                 bg_path = f"{self.output_dir}/bg_scene_{scene_num}.png"
@@ -641,10 +537,8 @@ class ProfessionalCartoonFilmGenerator:
                     print(f"   📁 Internal path: {bg_path}")
                     print(download_info)
                     
-                    # Clear GPU memory after each generation
-                    if self.device == "cuda":
-                        torch.cuda.empty_cache()
-                        gc.collect()
+                    # Clear memory after each generation
+                    gc.collect()
                 else:
                     print(f"❌ Failed to save background image: {bg_path}")
                 
@@ -667,6 +561,9 @@ class ProfessionalCartoonFilmGenerator:
         try:
             print("🎬 Setting up Open-Sora 2.0 for video generation...")
             
+            # Import torch here to avoid the UnboundLocalError
+            import torch
+            
             # Check available GPU memory
             if torch.cuda.is_available():
                 gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
@@ -774,7 +671,6 @@ class ProfessionalCartoonFilmGenerator:
             traceback.print_exc()
             return False
     
-    @spaces.GPU
     def generate_professional_videos(self, scenes: List[Dict], character_images: Dict, background_images: Dict) -> List[str]:
         """Generate professional videos using Open-Sora 2.0"""
         scene_videos = []
@@ -1129,11 +1025,10 @@ class ProfessionalCartoonFilmGenerator:
             return None
     
     def _create_emergency_fallback_video(self, script_data: Dict) -> str:
-        """Create a simple emergency fallback video when everything else fails"""
+        """Create emergency fallback video when all else fails"""
         try:
             print("🆘 Creating emergency fallback video...")
             
-            # Create a simple colored background
             width, height = 1024, 768
             background_color = (100, 150, 200)  # Blue-ish color
             
@@ -1152,12 +1047,16 @@ class ProfessionalCartoonFilmGenerator:
             
             # Create simple animated background
             for i in range(total_frames):
+                # Create frame with proper uint8 type
                 frame = np.full((height, width, 3), background_color, dtype=np.uint8)
                 
-                # Add simple animation (color shift)
+                # Add simple animation (color shift) with proper clamping
                 progress = i / total_frames
                 color_shift = int(50 * np.sin(progress * 2 * np.pi))
-                frame[:, :, 0] = np.clip(frame[:, :, 0] + color_shift, 0, 255)
+                
+                # Ensure all values stay within uint8 bounds (0-255)
+                new_blue = np.clip(frame[:, :, 0].astype(np.int16) + color_shift, 0, 255).astype(np.uint8)
+                frame[:, :, 0] = new_blue
                 
                 # Add text
                 font = cv2.FONT_HERSHEY_SIMPLEX
@@ -1181,6 +1080,8 @@ class ProfessionalCartoonFilmGenerator:
                 
         except Exception as e:
             print(f"❌ Emergency fallback video creation failed: {e}")
+            import traceback
+            traceback.print_exc()
             return None
     
     def merge_professional_film(self, scene_videos: List[str], script_data: Dict) -> str:
@@ -1224,12 +1125,25 @@ class ProfessionalCartoonFilmGenerator:
             print(f"❌ Video merging failed: {e}")
             return None
     
-    @spaces.GPU
     def generate_professional_cartoon_film(self, script: str) -> tuple:
-        """Main function to generate professional-quality cartoon film"""
+        """Main function to generate professional-quality cartoon film (ZeroGPU compatible)"""
         try:
             print("🎬 Starting professional cartoon film generation...")
             
+            # Step 0: Load models first (critical!)
+            print("🚀 Loading AI models...")
+            models_loaded = self.load_models()
+            if not models_loaded:
+                print("❌ Failed to load models - cannot generate content")
+                error_info = {
+                    "error": True,
+                    "message": "Failed to load AI models",
+                    "characters": [],
+                    "scenes": [],
+                    "style": "Model loading failed"
+                }
+                return None, error_info, "❌ Failed to load AI models", [], []
+            
             # Step 1: Generate professional script
             print("📝 Creating professional script structure...")
             script_data = self.generate_professional_script(script)
@@ -1486,9 +1400,8 @@ class ProfessionalCartoonFilmGenerator:
 # Initialize professional generator
 generator = ProfessionalCartoonFilmGenerator()
 
-@spaces.GPU
 def create_professional_cartoon_film(script):
-    """Gradio interface function for professional generation"""
+    """Gradio interface function for professional generation (ZeroGPU compatible)"""
     if not script.strip():
         empty_response = {
             "error": True,
@@ -1508,7 +1421,7 @@ def create_professional_cartoon_film(script):
             "scenes": [],
             "style": "Please wait for current generation to complete"
         }
-        return None, busy_response, "⏳ Another generation is in progress. Please wait and try again.", [], []
+        return None, busy_response, "⏳ Generation already in progress - please wait", [], []
     
     try:
         return generator.generate_professional_cartoon_film(script)