Create app.py

app.py

@@ -0,0 +1,321 @@
+import torch
+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 pathlib import Path
+import json
+import logging
+from dataclasses import dataclass
+import gc
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class GenerationConfig:
+    num_images: int = 1
+    num_inference_steps: int = 50
+    guidance_scale: float = 7.5
+    seed: Optional[int] = None
+
+class ModelCache:
+    def __init__(self, cache_dir: Path):
+        self.cache_dir = cache_dir
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def load_model(self, model_id: str, load_func: callable, cache_name: str) -> Any:
+        try:
+            logger.info(f"Loading {cache_name}")
+            return load_func(model_id)
+        except Exception as e:
+            logger.error(f"Error loading model {cache_name}: {str(e)}")
+            raise
+
+class EnhancedBanglaSDGenerator:
+    def __init__(
+        self,
+        banglaclip_weights_path: str,
+        cache_dir: str,
+        device: Optional[torch.device] = None
+    ):
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {self.device}")
+
+        self.cache = ModelCache(Path(cache_dir))
+        self._initialize_models(banglaclip_weights_path)
+        self._load_context_data()
+
+    def _initialize_models(self, banglaclip_weights_path: str):
+        try:
+            # Initialize translation models
+            self.bn2en_model_name = "Helsinki-NLP/opus-mt-bn-en"
+            self.translator = self.cache.load_model(
+                self.bn2en_model_name,
+                MarianMTModel.from_pretrained,
+                "translator"
+            ).to(self.device)
+            self.trans_tokenizer = MarianTokenizer.from_pretrained(self.bn2en_model_name)
+
+            # Initialize CLIP models
+            self.clip_model_name = "openai/clip-vit-base-patch32"
+            self.bangla_text_model = "csebuetnlp/banglabert"
+            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)
+
+            # Initialize Stable Diffusion
+            self._initialize_stable_diffusion()
+
+        except Exception as e:
+            logger.error(f"Error initializing models: {str(e)}")
+            raise RuntimeError(f"Failed to initialize models: {str(e)}")
+
+    def _initialize_stable_diffusion(self):
+        """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."""
+    cache_dir = Path("model_cache")
+    generator = None
+
+    def initialize_generator():
+        nonlocal generator
+        if generator is None:
+            generator = EnhancedBanglaSDGenerator(
+                banglaclip_weights_path="banglaclip_model_epoch_10.pth",
+                cache_dir=str(cache_dir)
+            )
+        return generator
+
+    def cleanup_generator():
+        nonlocal generator
+        if generator is not None:
+            generator.cleanup()
+            generator = None
+
+    def generate_images(text: str, num_images: int, steps: int, guidance_scale: float, seed: Optional[int]) -> Tuple[List[Any], str]:
+        if not text.strip():
+            return None, "দয়া করে কিছু টেক্সট লিখুন"
+
+        try:
+            gen = initialize_generator()
+            config = GenerationConfig(
+                num_images=int(num_images),
+                num_inference_steps=int(steps),
+                guidance_scale=float(guidance_scale),
+                seed=int(seed) if seed else None
+            )
+
+            images, prompt = gen.generate_image(text, config)
+            cleanup_generator()
+            return images, prompt
+
+        except Exception as e:
+            logger.error(f"Error in Gradio interface: {str(e)}")
+            cleanup_generator()
+            return None, f"ছবি তৈরি ব্যর্থ হয়েছে: {str(e)}"
+
+    # 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()
+    # Fixed queue configuration for newer Gradio versions
+    demo.queue().launch(share=True)