Text-to-Image Generation

class StableDiffusionPipeline:
    """
    Implementation of Stable Diffusion text-to-image generation
    Based on latent diffusion models architecture
    """
    
    def __init__(self, model_path, device='cuda'):
        self.device = device
        self.model_components = self.load_components(model_path)
        
    def load_components(self, model_path):
        """Load all Stable Diffusion components"""
        from diffusers import (
            AutoencoderKL, UNet2DConditionModel,
            PNDMScheduler, CLIPTextModel, CLIPTokenizer
        )
        
        components = {
            'vae': AutoencoderKL.from_pretrained(
                model_path, subfolder='vae'
            ).to(self.device),
            
            'unet': UNet2DConditionModel.from_pretrained(
                model_path, subfolder='unet'
            ).to(self.device),
            
            'scheduler': PNDMScheduler.from_pretrained(
                model_path, subfolder='scheduler'
            ),
            
            'text_encoder': CLIPTextModel.from_pretrained(
                model_path, subfolder='text_encoder'
            ).to(self.device),
            
            'tokenizer': CLIPTokenizer.from_pretrained(
                model_path, subfolder='tokenizer'
            )
        }
        
        return components
    
    def encode_prompt(self, prompt, negative_prompt=None):
        """Encode text prompt using CLIP"""
        import torch
        
        # Tokenize prompt
        text_inputs = self.model_components['tokenizer'](
            prompt,
            padding='max_length',
            max_length=77,
            truncation=True,
            return_tensors='pt'
        )
        
        # Get text embeddings
        with torch.no_grad():
            text_embeddings = self.model_components['text_encoder'](
                text_inputs.input_ids.to(self.device)
            ).last_hidden_state
        
        # Handle negative prompt
        if negative_prompt:
            negative_inputs = self.model_components['tokenizer'](
                negative_prompt,
                padding='max_length',
                max_length=77,
                truncation=True,
                return_tensors='pt'
            )
            
            with torch.no_grad():
                negative_embeddings = self.model_components['text_encoder'](
                    negative_inputs.input_ids.to(self.device)
                ).last_hidden_state
        else:
            # Use unconditional embeddings
            negative_embeddings = self.model_components['text_encoder'](
                self.model_components['tokenizer'](
                    '',
                    padding='max_length',
                    max_length=77,
                    return_tensors='pt'
                ).input_ids.to(self.device)
            ).last_hidden_state
        
        # Concatenate for classifier-free guidance
        text_embeddings = torch.cat([negative_embeddings, text_embeddings])
        
        return text_embeddings
    
    def diffusion_loop(self, latents, text_embeddings, num_steps=50, guidance_scale=7.5):
        """Main diffusion denoising loop"""
        import torch
        
        scheduler = self.model_components['scheduler']
        unet = self.model_components['unet']
        
        # Set timesteps
        scheduler.set_timesteps(num_steps)
        
        # Scale initial noise by scheduler
        latents = latents * scheduler.init_noise_sigma
        
        # Denoising loop
        for i, t in enumerate(scheduler.timesteps):
            # Expand latents for classifier-free guidance
            latent_model_input = torch.cat([latents] * 2)
            latent_model_input = scheduler.scale_model_input(latent_model_input, t)
            
            # Predict noise residual
            with torch.no_grad():
                noise_pred = unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=text_embeddings
                ).sample
            
            # Perform guidance
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (
                noise_pred_text - noise_pred_uncond
            )
            
            # Compute previous noisy sample
            latents = scheduler.step(noise_pred, t, latents).prev_sample
            
            # Optional: yield for progress tracking
            if i % 10 == 0:
                yield {'step': i, 'total': num_steps, 'latents': latents}
        
        return latents
    
    def decode_latents(self, latents):
        """Decode latents to image using VAE"""
        import torch
        
        vae = self.model_components['vae']
        
        # Scale latents
        latents = 1 / 0.18215 * latents
        
        # Decode to image
        with torch.no_grad():
            image = vae.decode(latents).sample
        
        # Convert to PIL Image
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
        image = (image * 255).astype('uint8')
        
        from PIL import Image
        return Image.fromarray(image[0])
    
    def generate(self, prompt, **kwargs):
        """Complete text-to-image generation pipeline"""
        import torch
        
        # Parameters
        height = kwargs.get('height', 512)
        width = kwargs.get('width', 512)
        num_steps = kwargs.get('num_inference_steps', 50)
        guidance_scale = kwargs.get('guidance_scale', 7.5)
        seed = kwargs.get('seed', None)
        
        # Set random seed for reproducibility
        if seed is not None:
            torch.manual_seed(seed)
        
        # Encode prompt
        text_embeddings = self.encode_prompt(
            prompt,
            kwargs.get('negative_prompt', None)
        )
        
        # Generate initial random noise
        latents = torch.randn(
            (1, 4, height // 8, width // 8),
            device=self.device
        )
        
        # Run diffusion
        for update in self.diffusion_loop(
            latents, text_embeddings, num_steps, guidance_scale
        ):
            if isinstance(update, dict):
                # Progress update
                continue
            else:
                latents = update
        
        # Decode to image
        image = self.decode_latents(latents)
        
        return image

# Advanced features
class AdvancedImageGeneration:
    """Advanced techniques for text-to-image generation"""
    
    def img2img_pipeline(self, init_image, prompt, strength=0.8):
        """Image-to-image generation with partial diffusion"""
        # Encode initial image to latents
        init_latents = self.encode_image(init_image)
        
        # Add noise based on strength
        num_inference_steps = 50
        start_step = int(num_inference_steps * (1 - strength))
        
        # Run partial diffusion from noised image
        noised_latents = self.add_noise(init_latents, start_step)
        final_latents = self.diffusion_loop(
            noised_latents, 
            text_embeddings,
            start_step=start_step
        )
        
        return self.decode_latents(final_latents)
    
    def inpainting_pipeline(self, image, mask, prompt):
        """Inpainting with masked diffusion"""
        # Encode image and mask
        image_latents = self.encode_image(image)
        mask_latents = self.encode_mask(mask)
        
        # Run masked diffusion
        for step in self.diffusion_steps:
            # Predict noise for masked region
            noise_pred = self.unet(masked_latents, step, text_embeddings)
            
            # Blend with original
            latents = mask_latents * noise_pred + (1 - mask_latents) * image_latents
        
        return self.decode_latents(latents)
    
    def controlnet_generation(self, prompt, control_image, control_type='canny'):
        """Controlled generation with ControlNet"""
        # Process control image
        if control_type == 'canny':
            control = self.extract_canny_edges(control_image)
        elif control_type == 'depth':
            control = self.estimate_depth(control_image)
        elif control_type == 'pose':
            control = self.detect_pose(control_image)
        
        # Run controlled diffusion
        latents = self.diffusion_with_control(
            text_embeddings,
            control_conditioning=control
        )
        
        return self.decode_latents(latents)