Update app.py

app.py

@@ -1,73 +1,172 @@
 import torch
 import gradio as gr
-from PIL import Image, ImageDraw
+import logging
+import traceback
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from diffusers import StableDiffusionPipeline
+from PIL import Image, ImageDraw
+
+# Logging Configuration
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')
+logger = logging.getLogger(__name__)
 
 class T4OptimizedStorylineGenerator:
     def __init__(self):
-        # Model Selection
-        text_model_name = "distilgpt2"
-        image_model_name = "runwayml/stable-diffusion-v1-5"
-
-        # Device Configuration
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
-
-        # Load Text Generation Model
-        self.tokenizer = AutoTokenizer.from_pretrained(text_model_name)
-        self.text_model = AutoModelForCausalLM.from_pretrained(
-            text_model_name,
-            torch_dtype=self.dtype,
-            device_map='auto'
-        ).to(self.device)
-
-        # Load Image Generation Pipeline
-        self.image_pipeline = StableDiffusionPipeline.from_pretrained(
-            image_model_name,
-            torch_dtype=self.dtype
-        )
-        if self.device == "cuda":
-            self.image_pipeline.enable_attention_slicing()
-        self.image_pipeline = self.image_pipeline.to(self.device)
+        """
+        Optimized initialization for both CPU and GPU environments
+        """
+        try:
+            # Model Selection
+            text_model_name = "distilgpt2"  # Lighter GPT-2 model
+            image_model_name = "runwayml/stable-diffusion-v1-5"  # Stable Diffusion model
+
+            # Device Configuration
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+            logger.info(f"🖥️ Device: {self.device}")
+            logger.info(f"📊 Precision: {self.dtype}")
+
+            # Text Generation Model
+            logger.info("🔄 Loading Text Model")
+            self.tokenizer = AutoTokenizer.from_pretrained(text_model_name)
+            self.text_model = AutoModelForCausalLM.from_pretrained(
+                text_model_name,
+                torch_dtype=self.dtype
+            ).to(self.device)  # Removed device_map to avoid accelerate dependency for CPU
+
+            # Image Generation Pipeline
+            logger.info("🖌️ Loading Image Generation Pipeline")
+            self.image_pipeline = StableDiffusionPipeline.from_pretrained(
+                image_model_name,
+                torch_dtype=self.dtype
+            )
+            
+            if self.device == "cuda":
+                self.image_pipeline.enable_attention_slicing()  # Memory optimization for GPU
+            self.image_pipeline = self.image_pipeline.to(self.device)
+
+        except Exception as e:
+            logger.error(f"Initialization Error: {e}")
+            logger.error(traceback.format_exc())
+            raise
 
     def generate_story_text(self, topic, plot_points, max_length=200):
-        prompt = f"Write a short story about {topic}. Scenes: {', '.join(plot_points)}"
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=100).to(self.device)
-        outputs = self.text_model.generate(
-            inputs.input_ids,
-            max_length=max_length,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_k=50,
-            top_p=0.95
-        )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        """
+        Optimized text generation for T4 with reduced complexity
+        """
+        try:
+            prompt = f"Write a short story about {topic}. Scenes: {', '.join(plot_points)}"
+            
+            # Tokenization with memory efficiency
+            inputs = self.tokenizer(
+                prompt, 
+                return_tensors="pt", 
+                truncation=True, 
+                max_length=100
+            ).to(self.device)
+
+            # Generation with reduced complexity
+            with torch.no_grad():
+                outputs = self.text_model.generate(
+                    inputs.input_ids,
+                    max_length=max_length,
+                    num_return_sequences=1,
+                    temperature=0.7,
+                    top_k=50,
+                    top_p=0.95
+                )
+            
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        except Exception as e:
+            logger.error(f"Text Generation Error: {e}")
+            return f"Story generation failed: {e}"
 
     def generate_scene_images(self, scene_prompts, num_images=3):
+        """
+        T4-Optimized image generation with reduced steps and memory usage
+        """
         scene_images = []
-        for prompt in scene_prompts[:num_images]:
-            image = self.image_pipeline(prompt, num_inference_steps=25, guidance_scale=6.0, height=512, width=512).images[0]
-            scene_images.append(image)
+        
+        try:
+            for prompt in scene_prompts[:num_images]:
+                with torch.inference_mode():
+                    # Reduced inference steps for T4
+                    image = self.image_pipeline(
+                        prompt, 
+                        num_inference_steps=25,  # Reduced from standard 50
+                        guidance_scale=6.0,
+                        height=512,   # Standard resolution
+                        width=512
+                    ).images[0]
+                
+                scene_images.append(image)
+        
+        except Exception as e:
+            logger.error(f"Image Generation Error: {e}")
+            # Fallback error image generation
+            scene_images.append(self._create_error_image())
+        
         return scene_images
 
+    def _create_error_image(self, size=(600, 400)):
+        """
+        Create a standardized error visualization image
+        """
+        img = Image.new('RGB', size, color=(200, 50, 50))
+        draw = ImageDraw.Draw(img)
+        
+        # Simple error message
+        draw.text(
+            (50, 180), 
+            "Image Generation Failed", 
+            fill=(255, 255, 255)
+        )
+        
+        return img
+
     def generate_visual_storyline(self, topic, plot_points, scene_prompts):
-        story_text = self.generate_story_text(topic, plot_points)
-        scene_images = self.generate_scene_images(scene_prompts)
-        return story_text, scene_images
+        """
+        Comprehensive storyline generation with T4 optimization
+        """
+        try:
+            # Generate story text
+            story_text = self.generate_story_text(topic, plot_points)
+            
+            # Generate scene images with T4 constraints
+            scene_images = self.generate_scene_images(scene_prompts)
+            
+            return story_text, scene_images
+        
+        except Exception as e:
+            logger.error(f"Visual Storyline Generation Error: {e}")
+            error_text = f"Storyline generation failed: {e}"
+            error_image = self._create_error_image()
+            
+            return error_text, [error_image]
 
 def create_gradio_interface():
+    """
+    Create a user-friendly Gradio interface with T4 optimizations
+    """
     generator = T4OptimizedStorylineGenerator()
 
     def storyline_wrapper(topic, plot_points, scene_prompts):
+        # Split input strings
         plot_points_list = [p.strip() for p in plot_points.split(',')]
         scene_prompts_list = [p.strip() for p in scene_prompts.split(',')]
-        return generator.generate_visual_storyline(topic, plot_points_list, scene_prompts_list)
+        
+        return generator.generate_visual_storyline(
+            topic, 
+            plot_points_list, 
+            scene_prompts_list
+        )
 
     interface = gr.Interface(
         fn=storyline_wrapper,
         inputs=[
-            gr.Textbox(label="Story Topic"),
+            gr.Textbox(label="Story Topic", placeholder="A magical adventure"),
             gr.Textbox(label="Plot Points", placeholder="Enter forest, Meet creatures, Find treasure"),
             gr.Textbox(label="Scene Prompts", placeholder="Misty enchanted forest, Magical creatures gathering, Hidden treasure cave")
         ],
@@ -75,11 +174,32 @@ def create_gradio_interface():
             gr.Textbox(label="Generated Story"),
             gr.Gallery(label="Scene Images")
         ],
-        title="T4-Optimized AI Storyline Generator"
+        title="🌈 T4-Optimized AI Storyline Generator",
+        description="Create magical stories with AI-powered text and image generation"
     )
     
     return interface
 
+def main():
+    """Main execution with error handling"""
+    try:
+        # Launch Gradio interface
+        interface = create_gradio_interface()
+        interface.launch(debug=True)
+    
+    except Exception as e:
+        logger.critical(f"Critical Failure: {e}")
+        logger.critical(traceback.format_exc())
+
 if __name__ == "__main__":
-    interface = create_gradio_interface()
-    interface.launch()
+    main()
+
+# T4 GPU Optimization Notes
+"""
+Optimization Strategies:
+1. Use distilgpt2 (lighter model)
+2. Reduced inference steps (25 vs 50)
+3. float16 precision when using GPU
+4. Attention slicing when using GPU
+5. Smaller image resolutions
+"""