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_quantized.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)