import torch
from diffusers import StableDiffusionPipeline

class StableDiffusionWrapper:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.pipe = self._load_model()
        
    def _load_model(self):
        """Load model directly from Hugging Face"""
        print("⏳ Loading Stable Diffusion...")
        pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            variant="fp16",
            safety_checker=None
        ).to(self.device)
        
        # Optimizations
        pipe.enable_attention_slicing()
        if torch.cuda.is_available():
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except:
                print("⚠️ XFormers not available")
        
        print("✅ Model ready")
        return pipe

    def generate(self, prompt, **kwargs):
        """Generate image from prompt"""
        return self.pipe(prompt, **kwargs).images[0]

# Singleton instance to be imported by Gradio app
sd_model = StableDiffusionWrapper()