from PIL import Image
import os
import ollama
import torch
from diffusers import DiffusionPipeline, AutoPipelineForImage2Image, LCMScheduler, AutoPipelineForText2Image
import time

model_id = "simianluo/lcm_dreamshaper_v7"

class GenerationSession:
    def __init__(self, model_id):
        self.model_id = model_id   
        self.txt2img_pipeline = None  
        self.img2img_pipeline = None
        self.current_image = None
        self.current_prompt = None
        self._initialize_pipelines()

    def _initialize_pipelines(self):
            print("initializing pipelines...")

            self.txt2img_pipeline = DiffusionPipeline.from_pretrained(
                model_id,
                torch_dtype = torch.float16,
                safety_checker = None
            )
            
            self.txt2img_pipeline.scheduler = LCMScheduler.from_config(self.txt2img_pipeline.scheduler.config)
            self.txt2img_pipeline.to("cuda")
            self.txt2img_pipeline.enable_attention_slicing()
            self.txt2img_pipeline.enable_vae_slicing()

           # self.txt2img_pipeline.unet = torch.compile(
           #   self.txt2img_pipeline.unet,
           #    mode = "reduce-overhead",
           #     fullgraph = True
           #)
            print("Text 2 image pipeline loaded and compiled.")


            self.img2img_pipeline = AutoPipelineForImage2Image.from_pipe(self.txt2img_pipeline)  
            print("Image 2 image pipeline loaded (shared weights).")

    def GeneratingBaseImage(self, prompt: str, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
        start = time.time()
        image = self.txt2img_pipeline(
            prompt = prompt,
            negative_prompt= negative_prompt,
            num_inference_steps = 4,
            guidance_scale = 1.0,
            height = 512,
            width = 512
        ).images
        print(f"Text to image generated in [{time.time() - start:.2f}s]")
        return image
    
    def GeneratingVariationImage(self, prompt: str, reference_image: Image.Image, strength: float = 0.5, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
        start = time.time()
        image = self.img2img_pipeline(
            prompt = prompt,
            image = reference_image,
            strength = strength,
            num_inference_steps = 4,
            guidance_scale = 1.0,
            negative_prompt = negative_prompt
        ).images
        print(f"Image to image generated in [{time.time() - start:.2f}s]")
        return image
    
    def Generate(self, new_prompt: str, strength: float = 0.5):
        if self.current_image is None:
            self.current_image = self.GeneratingBaseImage(new_prompt)
        else:
            self.current_image = self.GeneratingVariationImage(new_prompt, self.current_image, strength)
        
        self.current_prompt = new_prompt
        return self.current_image
    
    def reset(self):
        self.current_image = None
        self.current_prompt = None
        print("Session reset. Ready for new generation.")