import gradio as gr
import torch
from diffusers import FluxPipeline
import os

HF_TOKEN = os.environ.get("HF_TOKEN")

print("Loading FLUX.1-schnell model for CPU...")
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.float32,
    token=HF_TOKEN
)
pipe.to("cpu")
pipe.enable_attention_slicing(1)
print("Model loaded successfully on CPU")

def generate_image(prompt, width, height, num_steps, seed, progress=gr.Progress()):
    if not prompt:
        return None, "Masukkan prompt terlebih dahulu!"
    
    progress(0, desc="Memulai generate gambar...")
    
    if seed == -1:
        generator = None
    else:
        generator = torch.Generator("cpu").manual_seed(int(seed))
    
    try:
        progress(0.3, desc="Generating... Mohon tunggu 2-5 menit...")
        
        image = pipe(
            prompt=prompt,
            width=width,
            height=height,
            num_inference_steps=num_steps,
            guidance_scale=0.0,
            generator=generator,
            max_sequence_length=256
        ).images[0]
        
        progress(1.0, desc="Selesai!")
        return image, f"Gambar berhasil dibuat! ({width}x{height}, {num_steps} steps)"
        
    except Exception as e:
        return None, f"Error: {str(e)}"

with gr.Blocks(theme=gr.themes.Soft(), title="FLUX Image Generator") as demo:
    gr.Markdown(
        """
        # FLUX.1-schnell Image Generator
        Generate gambar berkualitas tinggi menggunakan AI model FLUX.1-schnell
        
        Running on CPU - Unlimited usage, no quota!
        Waktu generate: ~2-5 menit per gambar (tergantung ukuran)
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Contoh: A cute cat wearing sunglasses, digital art...",
                lines=4
            )
            
            with gr.Accordion("Pengaturan Lanjutan", open=False):
                with gr.Row():
                    width = gr.Slider(
                        minimum=256,
                        maximum=768,
                        value=512,
                        step=64,
                        label="Lebar (Width)"
                    )
                    height = gr.Slider(
                        minimum=256,
                        maximum=768,
                        value=512,
                        step=64,
                        label="Tinggi (Height)"
                    )
                
                num_steps = gr.Slider(
                    minimum=1,
                    maximum=8,
                    value=4,
                    step=1,
                    label="Inference Steps (4 = cepat & bagus)"
                )
                
                seed = gr.Number(
                    label="Seed (-1 untuk random)",
                    value=-1,
                    precision=0
                )
            
            generate_btn = gr.Button("Generate Gambar", variant="primary", size="lg")
            
            status = gr.Textbox(
                label="Status",
                value="Siap generate gambar...",
                interactive=False
            )
            
            gr.Markdown(
                """
                ### Tips:
                - **512x512**: Paling cepat (~2 menit)
                - **768x768**: Lebih detail (~4 menit)
                - **4 steps**: Balance terbaik speed & quality
                - **Prompt jelas**: Deskripsi detail = hasil lebih baik
                - **Seed**: Set angka spesifik untuk hasil konsisten
                """
            )
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Hasil Gambar",
                type="pil",
                height=550
            )
    
    gr.Examples(
        examples=[
            ["Seekor kucing lucu memakai kacamata hitam di pantai, digital art", 512, 512, 4, 42],
            ["Pemandangan gunung dengan sunset yang indah, highly detailed", 512, 512, 4, 123],
            ["Kota futuristik di malam hari dengan lampu neon, cyberpunk style", 512, 768, 4, 999],
            ["Portrait wanita cantik dengan bunga di rambut, oil painting", 512, 512, 4, 555],
            ["Naga besar terbang di atas kastil, fantasy art, epic", 768, 512, 4, 777],
            ["Astronaut riding a horse on mars, photorealistic, 4k", 512, 512, 4, 888],
        ],
        inputs=[prompt, width, height, num_steps, seed],
        outputs=[output_image, status],
        fn=generate_image,
        cache_examples=False,
        label="Contoh Prompt"
    )
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, width, height, num_steps, seed],
        outputs=[output_image, status]
    )
    
    gr.Markdown(
        """
        ---
        ### Informasi
        - Model: FLUX.1-schnell by Black Forest Labs
        - Hardware: CPU (Unlimited usage)
        - Framework: Diffusers + Gradio
        
        ### Kenapa CPU?
        - Unlimited - Tidak ada quota/limit
        - Gratis - 100% free forever
        - Stabil - Tidak ada antrian
        - Trade-off: Lebih lambat dari GPU
        """
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()