Update app.py

app.py

@@ -71,7 +71,7 @@ class EnhancedBanglaSDGenerator:
             self.processor = CLIPProcessor.from_pretrained(self.clip_model_name)
             self.tokenizer = AutoTokenizer.from_pretrained(self.bangla_text_model)
 
-            # Initialize Stable Diffusion
+            # Initialize Stable Diffusion with optimizations
             self._initialize_stable_diffusion()
 
         except Exception as e:
@@ -79,28 +79,38 @@ class EnhancedBanglaSDGenerator:
             raise RuntimeError(f"Failed to initialize models: {str(e)}")
 
     def _initialize_stable_diffusion(self):
-        """Initialize Stable Diffusion pipeline with optimized settings."""
+        """Initialize Stable Diffusion pipeline with CPU performance optimizations."""
         self.pipe = self.cache.load_model(
             "runwayml/stable-diffusion-v1-5",
             lambda model_id: StableDiffusionPipeline.from_pretrained(
                 model_id,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                safety_checker=None
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                use_safetensors=True,
+                use_memory_efficient_attention=True,
+                local_files_only=True
             ),
             "stable_diffusion"
         )
 
+        # Optimize scheduler
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
             self.pipe.scheduler.config,
             use_karras_sigmas=True,
             algorithm_type="dpmsolver++"
         )
-        self.pipe = self.pipe.to(self.device)
 
-        # Memory optimization
-        self.pipe.enable_attention_slicing()
-        if torch.cuda.is_available():
-            self.pipe.enable_sequential_cpu_offload()
+        # CPU optimizations
+        self.pipe.enable_attention_slicing(slice_size=1)
+        self.pipe.enable_vae_slicing()
+        self.pipe.enable_sequential_cpu_offload()
+        
+        # Component-level optimizations
+        for component in [self.pipe.text_encoder, self.pipe.vae, self.pipe.unet]:
+            if hasattr(component, 'enable_model_cpu_offload'):
+                component.enable_model_cpu_offload()
+        
+        self.pipe = self.pipe.to(self.device)
 
     def _load_banglaclip_model(self, weights_path: str) -> CLIPModel:
         try:
@@ -175,15 +185,27 @@ class EnhancedBanglaSDGenerator:
             enhanced_prompt = self._enhance_prompt(bangla_text)
             negative_prompt = self._get_negative_prompt()
 
-            with torch.autocast(self.device.type):
+            # Pre-generation optimization
+            torch.set_num_threads(max(4, torch.get_num_threads()))
+            gc.collect()
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+            # Memory-optimized generation
+            with torch.inference_mode():
                 result = self.pipe(
                     prompt=enhanced_prompt,
                     negative_prompt=negative_prompt,
                     num_images_per_prompt=config.num_images,
                     num_inference_steps=config.num_inference_steps,
-                    guidance_scale=config.guidance_scale
+                    guidance_scale=config.guidance_scale,
+                    use_memory_efficient_attention=True,
+                    use_memory_efficient_cross_attention=True
                 )
 
+            # Post-generation cleanup
+            gc.collect()
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
             return result.images, enhanced_prompt
 
         except Exception as e:
@@ -194,12 +216,10 @@ class EnhancedBanglaSDGenerator:
         """Enhance prompt with context and style information."""
         translated_text = self._translate_text(bangla_text)
 
-        # Gather contexts
         contexts = []
         contexts.extend(context for loc, context in self.location_contexts.items() if loc in bangla_text)
         contexts.extend(context for scene, context in self.scene_contexts.items() if scene in bangla_text)
 
-        # Add photo style
         photo_style = [
             "professional photography",
             "high resolution",
@@ -209,7 +229,6 @@ class EnhancedBanglaSDGenerator:
             "beautiful composition"
         ]
 
-        # Combine all parts
         all_parts = [translated_text] + contexts + photo_style
         return ", ".join(dict.fromkeys(all_parts))
 
@@ -319,5 +338,4 @@ def create_gradio_interface():
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    # Fixed queue configuration for newer Gradio versions
-    demo.queue().launch(share=True)
+    demo.queue().launch(share=True)