from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch
import gradio as gr
import spaces


lora_path = "OedoSoldier/detail-tweaker-lora"

@spaces.GPU
def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_scale=7.0,model="Real6.0"):
    """
    Generate an image using Stable Diffusion based on the input prompt
    """
        
    if model == "Real5.0":
        model_id = "SG161222/Realistic_Vision_V5.0_noVAE"
        
    elif model == "Real5.1":
        model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
        
    else:
        model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
        
        
    pipe = DiffusionPipeline.from_pretrained(model_id).to("cuda")
    
    if model == "Real6.0":
        pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))

    pipe.load_lora_weights(lora_path)

    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        algorithm_type="dpmsolver++",
        use_karras_sigmas=True
    )

    
    # Generate the image
    image = pipe(
        prompt = prompt,
        negative_prompt = negative_prompt,
        cross_attention_kwargs = {"scale":1},
        num_inference_steps = num_inference_steps,
        guidance_scale = guidance_scale,
        width = 960,
        height = 960
    ).images[0]
    
    return image

title = """<h1 align="center">ProFaker</h1>"""
# Create the Gradio interface
with gr.Blocks() as demo:
    gr.HTML(title)
    
    with gr.Row():
        with gr.Column():
            # Input components
            prompt = gr.Textbox(
                label="Prompt",
                info="Enter your image description here...",
                lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                info="Enter what you don't want in Image...",
                lines=3
            )
            generate_button = gr.Button("Generate Image")
            with gr.Accordion("Advanced Options", open=False):
                model = gr.Dropdown(
                    choices=["Real6.0","Real5.1","Real5.0"],
                    value="Real6.0",
                    label="Model",
                )
                steps_slider = gr.Slider(
                    minimum=1,
                    maximum=100,
                    value=30,
                    step=1,
                    label="Number of Steps"
                )
                guidance_slider = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=7.0,
                    step=0.5,
                    label="Guidance Scale"
                )
        with gr.Column():
            # Output component
            image_output = gr.Image(label="Generated Image")
            
    
    # Connect the interface to the generation function
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps_slider, guidance_slider, model],
        outputs=image_output
    )

demo.queue(max_size=10).launch(share=False)