import gradio as gr
import torch
from diffusers import StableDiffusionPipeline  # ← Ganti ke StableDiffusionPipeline
import gc
import time
import os

print("=" * 50)
print("🚀 Starting Image Generator on Hugging Face Free Tier")
print("=" * 50)

# MODEL YANG PASTI BISA dan MASIH AKTIF
MODEL_ID = "CompVis/stable-diffusion-v1-4"

print(f"📦 Loading model: {MODEL_ID}")
print("⚠️ Download model ~1.7GB, butuh 3-5 menit... SABAR YA!")

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"💻 Device: {device}")

# Clean memory
gc.collect()
if device == "cuda":
    torch.cuda.empty_cache()

# Load model
try:
    start_time = time.time()
    
    # Gunakan StableDiffusionPipeline, bukan DiffusionPipeline
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_safetensors=True,
        safety_checker=None,  # Matikan safety checker untuk hemat memory
        requires_safety_checker=False
    )
    
    # Enable memory optimizations
    pipe.enable_attention_slicing()
    
    if device == "cuda":
        pipe = pipe.to(device)
        pipe.enable_model_cpu_offload()
    else:
        print("⚠️ Running on CPU - akan lambat tapi stabil")
        pipe.enable_vae_slicing()
    
    load_time = time.time() - start_time
    print(f"✅ Model loaded in {load_time:.1f} seconds")
    print("🎉 Aplikasi siap digunakan!")
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    print("💡 Cek LOGS di atas untuk detail error")
    raise

def generate(prompt, negative_prompt, steps, guidance, width, height, seed):
    """Generate image with optimal settings for free tier"""
    
    print(f"\n🎨 Generating image...")
    print(f"📝 Prompt: {prompt[:50]}...")
    print(f"⚙️ Settings: steps={steps}, guidance={guidance}, {width}x{height}")
    
    start_time = time.time()
    
    try:
        # Pastikan resolusi multiple dari 8 (requirement SD)
        width = (width // 8) * 8
        height = (height // 8) * 8
        
        # Batasi resolusi untuk free tier CPU
        if device == "cpu":
            if width > 640:
                width = 640
            if height > 640:
                height = 640
        
        generator = torch.manual_seed(seed) if seed != -1 else None
        
        with torch.no_grad():
            image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                guidance_scale=guidance,
                width=width,
                height=height,
                generator=generator
            ).images[0]
        
        # Clean up memory
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()
        
        gen_time = time.time() - start_time
        print(f"✅ Generated in {gen_time:.1f} seconds")
        
        return image
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None

# Create the interface
with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Generator Free") as demo:
    gr.Markdown("""
    # 🎨 AI Image Generator - Free Tier
    
    **Model:** Stable Diffusion 1.4 (CompVis) ✅ Public & Active
    **Status:** Siap digunakan!
    
    ⚡ **Tips untuk hasil terbaik:**
    - Gunakan resolusi **512x512** untuk kecepatan optimal di CPU
    - **Steps 15-20** adalah sweet spot (cepat + quality cukup)
    - Hasil akan muncul dalam **1-3 menit** (CPU Free Tier)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="📝 Prompt", 
                placeholder="Describe the image you want to generate...",
                lines=3,
                value="a beautiful sunset over mountains, digital art, highly detailed"
            )
            
            negative_prompt = gr.Textbox(
                label="❌ Negative Prompt", 
                value="blurry, low quality, bad anatomy, distorted, ugly, worst quality, deformed",
                lines=2
            )
            
            with gr.Row():
                steps = gr.Slider(
                    10, 30, value=18, step=1, 
                    label="🔄 Steps (15-20 recommended)"
                )
                guidance = gr.Slider(
                    1, 15, value=7.0, step=0.5, 
                    label="🎯 Guidance Scale"
                )
            
            with gr.Row():
                width = gr.Dropdown(
                    [384, 512, 640],  # Hapus 768 karena terlalu berat untuk CPU
                    value=512, 
                    label="📐 Width"
                )
                height = gr.Dropdown(
                    [384, 512, 640], 
                    value=512, 
                    label="📐 Height"
                )
            
            seed = gr.Number(
                value=-1, 
                label="🎲 Seed (-1 for random)", 
                precision=0
            )
            
            btn = gr.Button("🚀 Generate Image", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output = gr.Image(label="✨ Generated Image", height=450)
    
    # Examples
    gr.Examples(
        examples=[
            ["a beautiful sunset over mountains, digital art, vibrant colors", "blurry, low quality", 18, 7.0, 512, 512, -1],
            ["a majestic lion in the savanna, wildlife photography", "blurry, low quality", 20, 7.5, 512, 512, -1],
            ["a futuristic city with flying cars, cyberpunk style", "blurry, ugly, low quality", 20, 7.5, 512, 512, 42],
            ["a cute puppy playing in grass, photorealistic", "cartoon, anime, blurry", 15, 7.0, 512, 512, -1],
        ],
        inputs=[prompt, negative_prompt, steps, guidance, width, height, seed],
        label="Click any example to test 👆"
    )
    
    gr.Markdown("""
    ---
    **ℹ️ Info:** 
    - ✅ **Model:** CompVis/stable-diffusion-v1-4 (Public, No Login)
    - 🚀 Loading model: 2-4 menit (first time only)
    - ⏱️ Generate image: 1-3 menit (CPU Free Tier)
    - 💾 Pastikan requirements.txt sudah sesuai
    """)

if __name__ == "__main__":
    print("🚀 Launching Gradio app...")
    demo.launch(server_name="0.0.0.0", server_port=7860)