import spaces
import gradio as gr
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler
from huggingface_hub import hf_hub_download
from PIL import Image
import requests
from translatepy import Translator
import numpy as np
import random
import os
hf_token = os.environ.get('HF_TOKEN')
from io import BytesIO

translator = Translator()

# Constants
model = "black-forest-labs/FLUX.1-dev"



MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

# Ensure model and scheduler are initialized in GPU-enabled function
if torch.cuda.is_available():
    transformer = FluxTransformer2DModel.from_single_file(
        "https://huggingface.co/ekt1701/Test_case/blob/main/rayflux_photoplus.safetensors",
        torch_dtype=torch.bfloat16
    )
    pipe = FluxPipeline.from_pretrained(
        model, 
        transformer=transformer,
        torch_dtype=torch.bfloat16, token=hf_token)
    pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
        pipe.scheduler.config, use_beta_sigmas=True
    )
    pipe.to("cuda")
    
    
@spaces.GPU()
def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
            prompt = prompt, 
            width = width,
            height = height,
            num_inference_steps = num_inference_steps,
            guidance_scale=guidance_scale,
            num_images_per_prompt=nums,
            generator = generator
    ).images

    
    return image, seed


css="""
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.HTML("<h1><center>Flux Playground</center></h1><p><center>RayFlux PhotoPlus Model.</center></p>")
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=2,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Gallery(label="Gallery", format="png", columns = 1, preview=True, height=400)
        
        with gr.Accordion("Advanced Settings", open=False):
                       
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=30,
                )

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=0,
                    maximum=10,
                    step=0.1,
                    value=3.5,
                )


            with gr.Row():
            
                nums = gr.Slider(
                    label="Number of Images",
                    minimum=1,
                    maximum=2,
                    step=1,
                    value=1,
                    scale=1,
                )
       
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=-1,
                )
            
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
        outputs = [result, seed]
    )
    
demo.launch()