Update app.py

app.py

@@ -1,13 +1,17 @@
+!wget -O banglaclip_model_epoch_10.pth https://huggingface.co/Mansuba/BanglaCLIP13/resolve/main/banglaclip_model_epoch_10.pth
+
+     
 import torch
-import os
-import requests
-import logging
-from pathlib import Path
 from transformers import CLIPModel, CLIPProcessor, AutoTokenizer, MarianMTModel, MarianTokenizer
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+import numpy as np
+from typing import List, Tuple, Optional, Dict, Any
 import gradio as gr
-from typing import List, Tuple, Optional, Any
+from pathlib import Path
+import json
+import logging
 from dataclasses import dataclass
+import gc
 
 # Configure logging
 logging.basicConfig(
@@ -16,30 +20,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-def download_model(model_url: str, model_path: str):
-    """Download large model file with progress tracking."""
-    if not os.path.exists(model_path):
-        try:
-            logger.info(f"Downloading model from {model_url}...")
-            response = requests.get(model_url, stream=True)
-            response.raise_for_status()
-
-            total_size = int(response.headers.get('content-length', 0))
-            block_size = 1024 * 1024  # 1 MB chunks
-            downloaded_size = 0
-
-            with open(model_path, 'wb') as f:
-                for data in response.iter_content(block_size):
-                    f.write(data)
-                    downloaded_size += len(data)
-                    progress = (downloaded_size / total_size) * 100 if total_size > 0 else 0
-                    logger.info(f"Download progress: {progress:.2f}%")
-
-            logger.info("Model download complete.")
-        except Exception as e:
-            logger.error(f"Model download failed: {e}")
-            raise
-
 @dataclass
 class GenerationConfig:
     num_images: int = 1
@@ -63,6 +43,7 @@ class ModelCache:
 class EnhancedBanglaSDGenerator:
     def __init__(
         self,
+        banglaclip_weights_path: str,
         cache_dir: str,
         device: Optional[torch.device] = None
     ):
@@ -70,21 +51,12 @@ class EnhancedBanglaSDGenerator:
         logger.info(f"Using device: {self.device}")
 
         self.cache = ModelCache(Path(cache_dir))
-        self._initialize_models()
+        self._initialize_models(banglaclip_weights_path)
         self._load_context_data()
 
-    def _load_banglaclip_model(self):
-        """Load BanglaCLIP model from Hugging Face directly"""
-        try:
-            model = CLIPModel.from_pretrained("Mansuba/BanglaCLIP13")
-            return model.to(self.device)
-        except Exception as e:
-            logger.error(f"Failed to load BanglaCLIP model: {str(e)}")
-            raise
-
-    def _initialize_models(self):
+    def _initialize_models(self, banglaclip_weights_path: str):
         try:
-            # Translation models
+            # Initialize translation models
             self.bn2en_model_name = "Helsinki-NLP/opus-mt-bn-en"
             self.translator = self.cache.load_model(
                 self.bn2en_model_name,
@@ -93,17 +65,14 @@ class EnhancedBanglaSDGenerator:
             ).to(self.device)
             self.trans_tokenizer = MarianTokenizer.from_pretrained(self.bn2en_model_name)
 
-            # CLIP models
+            # Initialize CLIP models
             self.clip_model_name = "openai/clip-vit-base-patch32"
             self.bangla_text_model = "csebuetnlp/banglabert"
-            
-            # Load BanglaCLIP model directly from Hugging Face
-            self.banglaclip_model = self._load_banglaclip_model()
-            
+            self.banglaclip_model = self._load_banglaclip_model(banglaclip_weights_path)
             self.processor = CLIPProcessor.from_pretrained(self.clip_model_name)
             self.tokenizer = AutoTokenizer.from_pretrained(self.bangla_text_model)
 
-            # Stable Diffusion
+            # Initialize Stable Diffusion
             self._initialize_stable_diffusion()
 
         except Exception as e:
@@ -111,8 +80,156 @@ class EnhancedBanglaSDGenerator:
             raise RuntimeError(f"Failed to initialize models: {str(e)}")
 
     def _initialize_stable_diffusion(self):
-        """Load and initialize Stable Diffusion pipeline"""
-        pass  # Your existing code for initializing Stable Diffusion
+        """Initialize Stable Diffusion pipeline with optimized settings."""
+        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
+            ),
+            "stable_diffusion"
+        )
+
+        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()
+
+    def _load_banglaclip_model(self, weights_path: str) -> CLIPModel:
+        try:
+            if not Path(weights_path).exists():
+                raise FileNotFoundError(f"BanglaCLIP weights not found at {weights_path}")
+
+            clip_model = CLIPModel.from_pretrained(self.clip_model_name)
+            state_dict = torch.load(weights_path, map_location=self.device)
+
+            cleaned_state_dict = {
+                k.replace('module.', '').replace('clip.', ''): v
+                for k, v in state_dict.items()
+                if k.replace('module.', '').replace('clip.', '').startswith(('text_model.', 'vision_model.'))
+            }
+
+            clip_model.load_state_dict(cleaned_state_dict, strict=False)
+            return clip_model.to(self.device)
+
+        except Exception as e:
+            logger.error(f"Failed to load BanglaCLIP model: {str(e)}")
+            raise
+
+    def _load_context_data(self):
+        """Load location and scene context data."""
+        self.location_contexts = {
+            'কক্সবাজার': 'Cox\'s Bazar beach, longest natural sea beach in the world, sandy beach',
+            'সেন্টমার্টিন': 'Saint Martin\'s Island, coral island, tropical paradise',
+            'সুন্দরবন': 'Sundarbans mangrove forest, Bengal tigers, riverine forest'
+        }
+
+        self.scene_contexts = {
+            'সৈকত': 'beach, seaside, waves, sandy shore, ocean view',
+            'সমুদ্র': 'ocean, sea waves, deep blue water, horizon',
+            'পাহাড়': 'mountains, hills, valleys, scenic landscape'
+        }
+
+    def _translate_text(self, bangla_text: str) -> str:
+        """Translate Bangla text to English."""
+        inputs = self.trans_tokenizer(bangla_text, return_tensors="pt", padding=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.translator.generate(**inputs)
+
+        translated = self.trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return translated
+
+    def _get_text_embedding(self, text: str):
+        """Get text embedding from BanglaCLIP model."""
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.banglaclip_model.get_text_features(**inputs)
+
+        return outputs
+
+    def generate_image(
+        self,
+        bangla_text: str,
+        config: Optional[GenerationConfig] = None
+    ) -> Tuple[List[Any], str]:
+        if not bangla_text.strip():
+            raise ValueError("Empty input text")
+
+        config = config or GenerationConfig()
+
+        try:
+            if config.seed is not None:
+                torch.manual_seed(config.seed)
+
+            enhanced_prompt = self._enhance_prompt(bangla_text)
+            negative_prompt = self._get_negative_prompt()
+
+            with torch.autocast(self.device.type):
+                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
+                )
+
+            return result.images, enhanced_prompt
+
+        except Exception as e:
+            logger.error(f"Error during image generation: {str(e)}")
+            raise
+
+    def _enhance_prompt(self, bangla_text: str) -> str:
+        """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",
+            "4k",
+            "detailed",
+            "realistic",
+            "beautiful composition"
+        ]
+
+        # Combine all parts
+        all_parts = [translated_text] + contexts + photo_style
+        return ", ".join(dict.fromkeys(all_parts))
+
+    def _get_negative_prompt(self) -> str:
+        return (
+            "blurry, low quality, pixelated, cartoon, anime, illustration, "
+            "painting, drawing, artificial, fake, oversaturated, undersaturated"
+        )
+
+    def cleanup(self):
+        """Clean up GPU memory"""
+        if hasattr(self, 'pipe'):
+            del self.pipe
+        if hasattr(self, 'banglaclip_model'):
+            del self.banglaclip_model
+        if hasattr(self, 'translator'):
+            del self.translator
+        torch.cuda.empty_cache()
+        gc.collect()
 
 def create_gradio_interface():
     """Create and configure the Gradio interface."""
@@ -123,6 +240,7 @@ def create_gradio_interface():
         nonlocal generator
         if generator is None:
             generator = EnhancedBanglaSDGenerator(
+                banglaclip_weights_path="banglaclip_model_epoch_10.pth",
                 cache_dir=str(cache_dir)
             )
         return generator
@@ -155,20 +273,53 @@ def create_gradio_interface():
             cleanup_generator()
             return None, f"ছবি তৈরি ব্যর্থ হয়েছে: {str(e)}"
 
-    with gr.Blocks() as demo:
-        text_input = gr.Textbox(label="Text", placeholder="Enter your prompt here...")
-        num_images_input = gr.Slider(minimum=1, maximum=5, value=1, label="Number of Images")
-        steps_input = gr.Slider(minimum=1, maximum=100, value=50, label="Steps")
-        guidance_scale_input = gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale")
-        seed_input = gr.Number(label="Seed", optional=True)
-
-        output_images = gr.Gallery(label="Generated Images")
-
-        generate_button = gr.Button("Generate Images")
-        generate_button.click(generate_images, inputs=[text_input, num_images_input, steps_input, guidance_scale_input, seed_input], outputs=[output_images])
+    # Create Gradio interface
+    demo = gr.Interface(
+        fn=generate_images,
+        inputs=[
+            gr.Textbox(
+                label="বাংলা টেক্সট লিখুন",
+                placeholder="যেকোনো বাংলা টেক্সট লিখুন...",
+                lines=3
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=1,
+                label="ছবির সংখ্যা"
+            ),
+            gr.Slider(
+                minimum=20,
+                maximum=100,
+                step=1,
+                value=50,
+                label="স্টেপস"
+            ),
+            gr.Slider(
+                minimum=1.0,
+                maximum=20.0,
+                step=0.5,
+                value=7.5,
+                label="গাইডেন্স স্কেল"
+            ),
+            gr.Number(
+                label="সীড (ঐচ্ছিক)",
+                precision=0
+            )
+        ],
+        outputs=[
+            gr.Gallery(label="তৈরি করা ছবি"),
+            gr.Textbox(label="ব্যবহৃত প্রম্পট")
+        ],
+        title="বাংলা টেক্সট থেকে ছবি তৈরি",
+        description="যেকোনো বাংলা টেক্সট দিয়ে উচ্চমানের ছবি তৈরি করুন"
+    )
 
     return demo
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    demo.queue().launch(share=True, debug=True)
+    # Fixed queue configuration for newer Gradio versions
+    demo.queue().launch(share=True)
+