app.py

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
import os

# 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)
        
        # Set environment variables for better memory management
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
        os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

    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}")
        
        # Set memory split for VRAM usage on CPU
        self.memory_split = 0.5  # Use 50% of available VRAM
        self.setup_memory_management()
        
        self.cache = ModelCache(Path(cache_dir))
        self._initialize_models(banglaclip_weights_path)
        self._load_context_data()

    def setup_memory_management(self):
        """Setup optimal memory management for CPU and VRAM"""
        if torch.cuda.is_available():
            total_memory = torch.cuda.get_device_properties(0).total_memory
            torch.cuda.set_per_process_memory_fraction(self.memory_split)
        
        # Optimize CPU memory
        torch.set_num_threads(min(8, os.cpu_count() or 4))
        torch.set_num_interop_threads(min(8, os.cpu_count() or 4))

    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,
                lambda x: MarianMTModel.from_pretrained(x, low_cpu_mem_usage=True),
                "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."""
        try:
            self.pipe = self.cache.load_model(
                "runwayml/stable-diffusion-v1-5",
                lambda model_id: StableDiffusionPipeline.from_pretrained(
                    model_id,
                    torch_dtype=torch.float32,
                    safety_checker=None,
                    use_safetensors=True,
                    low_cpu_mem_usage=True,
                ),
                "stable_diffusion"
            )

            # Optimize scheduler for speed
            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
                self.pipe.scheduler.config,
                use_karras_sigmas=True,
                algorithm_type="dpmsolver++",
                solver_order=2
            )

            # Memory optimizations
            self.pipe.enable_attention_slicing(slice_size=1)
            self.pipe.enable_vae_slicing()
            self.pipe.enable_sequential_cpu_offload()
            
            # VRAM optimization
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                self.pipe.enable_model_cpu_offload()
            
            self.pipe = self.pipe.to(self.device)
            
        except Exception as e:
            logger.error(f"Error initializing Stable Diffusion: {str(e)}")
            raise

    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,
                low_cpu_mem_usage=True
            )
            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(), torch.cpu.amp.autocast():
            outputs = self.translator.generate(**inputs)

        translated = self.trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
        return translated

    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)

            # Clear memory before generation
            gc.collect()
            torch.cuda.empty_cache() if torch.cuda.is_available() else None

            enhanced_prompt = self._enhance_prompt(bangla_text)
            negative_prompt = self._get_negative_prompt()

            # Use mixed precision for faster generation
            with torch.inference_mode(), torch.cpu.amp.autocast():
                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,
                    use_memory_efficient_attention=True,
                    use_memory_efficient_cross_attention=True
                )

            # Clear memory after generation
            gc.collect()
            torch.cuda.empty_cache() if torch.cuda.is_available() else None

            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)

        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)

        photo_style = [
            "professional photography",
            "high resolution",
            "4k",
            "detailed",
            "realistic",
            "beautiful composition"
        ]

        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__":
    # Set environment variables for better performance
    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
    
    demo = create_gradio_interface()
    demo.queue().launch(share=True)