app.py

# ===== CRITICAL: Import spaces FIRST =====
try:
    import spaces
    HF_SPACES = True
except ImportError:
    def spaces_gpu_decorator(duration=60):
        def decorator(func):
            return func
        return decorator
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()
    HF_SPACES = False

# ===== Core imports =====
import random
import os
import gradio as gr
import torch
from diffusers import DiffusionPipeline
from PIL import Image
import gc

# ===== Configuration =====
class Config:
    MODEL_ID = "stabilityai/sdxl-turbo"
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
    MAX_SEED = 2**32 - 1

# ===== Adult Content Styles =====
ADULT_STYLES = {
    "Natural": "natural beauty, soft lighting, intimate",
    "Artistic": "artistic nude, fine art photography, elegant",
    "Sensual": "sensual pose, romantic lighting, beautiful",
    "Glamour": "glamour photography, professional, stunning",
    "Vintage": "vintage pin-up style, retro aesthetic",
    "Boudoir": "boudoir photography, intimate setting, elegant"
}

# ===== Adult Prompt Templates =====
ADULT_PROMPTS = {
    "Female": [
        "beautiful woman, natural pose, soft lighting, artistic",
        "elegant female model, professional photography, stunning",
        "gorgeous woman, intimate setting, romantic lighting",
        "beautiful lady, sensual pose, artistic photography",
        "stunning female, glamour shot, professional lighting"
    ],
    "Male": [
        "handsome man, artistic pose, dramatic lighting",
        "attractive male model, professional photography",
        "muscular man, artistic lighting, confident pose",
        "good-looking guy, intimate setting, soft lighting"
    ],
    "Couple": [
        "romantic couple, intimate moment, soft lighting",
        "beautiful pair, sensual pose, artistic photography",
        "lovers embrace, romantic setting, warm lighting"
    ]
}

# ===== Pipeline Manager =====
class FastPipeline:
    def __init__(self):
        self.pipe = None
        self.loaded = False
    
    def load(self):
        if self.loaded:
            return True
        
        try:
            print("Loading SDXL-Turbo...")
            self.pipe = DiffusionPipeline.from_pretrained(
                Config.MODEL_ID,
                torch_dtype=Config.TORCH_DTYPE,
                use_safetensors=True,
                variant="fp16" if Config.TORCH_DTYPE == torch.float16 else None
            )
            self.pipe.to(Config.DEVICE)
            
            # Speed optimizations
            if hasattr(self.pipe, 'enable_attention_slicing'):
                self.pipe.enable_attention_slicing()
            if hasattr(self.pipe, 'enable_vae_slicing'):
                self.pipe.enable_vae_slicing()
            
            self.loaded = True
            print("✅ Model loaded successfully!")
            return True
        except Exception as e:
            print(f"❌ Failed to load model: {e}")
            return False
    
    def generate(self, prompt, negative_prompt="", width=512, height=512, steps=2, guidance=0.0, seed=None):
        if not self.loaded:
            if not self.load():
                return None
        
        if seed is None:
            seed = random.randint(0, Config.MAX_SEED)
        
        generator = torch.Generator(Config.DEVICE).manual_seed(seed)
        
        try:
            # Clear cache
            if Config.DEVICE == "cuda":
                torch.cuda.empty_cache()
            
            # Generate (ultra-fast settings)
            result = self.pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=steps,
                guidance_scale=guidance,
                generator=generator
            ).images[0]
            
            return result, seed
        except Exception as e:
            print(f"Generation failed: {e}")
            if Config.DEVICE == "cuda":
                torch.cuda.empty_cache()
            return None, seed

# ===== Global Pipeline =====
pipeline = FastPipeline()

# ===== Main Generation Function =====
@spaces.GPU(duration=20)
def generate_adult_image(
    prompt_type,
    custom_prompt,
    style,
    negative_prompt,
    width,
    height,
    steps,
    seed,
    randomize_seed
):
    try:
        # Get prompt
        if custom_prompt.strip():
            base_prompt = custom_prompt
        else:
            prompts = ADULT_PROMPTS.get(prompt_type, ADULT_PROMPTS["Female"])
            base_prompt = random.choice(prompts)
        
        # Add style
        style_text = ADULT_STYLES.get(style, "")
        final_prompt = f"{base_prompt}, {style_text}" if style_text else base_prompt
        
        # Generate seed
        if randomize_seed:
            seed = random.randint(0, Config.MAX_SEED)
        
        # Generate image
        result = pipeline.generate(
            prompt=final_prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=steps,
            seed=seed
        )
        
        if result[0] is None:
            return None, seed, "❌ Generation failed"
        
        image, used_seed = result
        info = f"✅ Generated!\nSeed: {used_seed}\nPrompt: {final_prompt}"
        
        return image, used_seed, info
        
    except Exception as e:
        return None, seed, f"❌ Error: {str(e)}"

# ===== Quick Prompt Generator =====
def get_random_prompt(prompt_type):
    prompts = ADULT_PROMPTS.get(prompt_type, ADULT_PROMPTS["Female"])
    return random.choice(prompts)

# ===== Interface =====
def create_interface():
    with gr.Blocks(title="Adult Content Generator") as demo:
        
        gr.HTML("""
        <div style="text-align: center; margin-bottom: 2rem;">
            <h1 style="color: #FF6B6B;">🔥 Image Content Generator</h1>
            <p style="color: #666;">Fast image generation with SDXL-Turbo</p>
            <p style="color: #FF4444; font-weight: bold;"></p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Quick prompts
                prompt_type = gr.Dropdown(
                    choices=list(ADULT_PROMPTS.keys()),
                    value="Female",
                    label="Quick Prompts"
                )
                
                random_btn = gr.Button("🎲 Random Prompt", variant="secondary")
                
                # Custom prompt
                custom_prompt = gr.Textbox(
                    label="Custom Prompt (optional)",
                    placeholder="Enter your own prompt...",
                    lines=3
                )
                
                # Style
                style = gr.Dropdown(
                    choices=list(ADULT_STYLES.keys()),
                    value="Natural",
                    label="Style"
                )
                
                # Settings
                with gr.Row():
                    width = gr.Slider(256, 1024, 512, step=64, label="Width")
                    height = gr.Slider(256, 1024, 512, step=64, label="Height")
                
                with gr.Row():
                    steps = gr.Slider(1, 8, 2, step=1, label="Steps (1-2 for speed)")
                    seed = gr.Number(42, label="Seed", precision=0)
                
                randomize_seed = gr.Checkbox(True, label="Random Seed")
                
                # Advanced
                with gr.Accordion("Advanced", open=False):
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        value="ugly, deformed, blurry, bad anatomy, worst quality, low quality",
                        lines=2
                    )
                
                # Generate button
                generate_btn = gr.Button(
                    "🔥 Generate Image Content",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # Result
                result_image = gr.Image(
                    label="Generated Image",
                    height=500
                )
                
                # Info
                info_text = gr.Textbox(
                    label="Generation Info",
                    lines=4,
                    interactive=False
                )
        
        # Quick presets
        with gr.Row():
            preset_natural = gr.Button("🌸 Natural Beauty")
            preset_artistic = gr.Button("🎨 Artistic Nude")
            preset_glamour = gr.Button("✨ Glamour Shot")
            preset_vintage = gr.Button("📸 Vintage Pin-up")
        
        # Event handlers
        def apply_preset(style_name, prompt_type_val="Female"):
            prompt = get_random_prompt(prompt_type_val)
            return prompt, style_name
        
        # Connect events
        random_btn.click(
            fn=get_random_prompt,
            inputs=[prompt_type],
            outputs=[custom_prompt]
        )
        
        generate_btn.click(
            fn=generate_adult_image,
            inputs=[
                prompt_type, custom_prompt, style, negative_prompt,
                width, height, steps, seed, randomize_seed
            ],
            outputs=[result_image, seed, info_text]
        )
        
        # Quick presets
        preset_natural.click(
            lambda: apply_preset("Natural"),
            outputs=[custom_prompt, style]
        )
        
        preset_artistic.click(
            lambda: apply_preset("Artistic"),
            outputs=[custom_prompt, style]
        )
        
        preset_glamour.click(
            lambda: apply_preset("Glamour"),
            outputs=[custom_prompt, style]
        )
        
        preset_vintage.click(
            lambda: apply_preset("Vintage"),
            outputs=[custom_prompt, style]
        )
        
        # Enter key support
        custom_prompt.submit(
            fn=generate_adult_image,
            inputs=[
                prompt_type, custom_prompt, style, negative_prompt,
                width, height, steps, seed, randomize_seed
            ],
            outputs=[result_image, seed, info_text]
        )
    
    return demo

# ===== Speed Optimizations =====
def optimize_speed():
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.cuda.empty_cache()

# ===== Main =====
def main():
    print("🔥 Image Content Generator with SDXL-Turbo")
    print("Image Content")
    
    # Optimize for speed
    optimize_speed()
    
    # Create interface
    demo = create_interface()
    
    # Launch
    launch_kwargs = {
        "server_name": "0.0.0.0",
        "server_port": 7860,
        "share": HF_SPACES,
        "show_error": True
    }
    
    demo.launch(**launch_kwargs)

if __name__ == "__main__":
    main()