import gradio as gr
import spaces

@spaces.GPU
def generate_image(prompt, negative_prompt="", seed=42, width=512, height=512, guidance_scale=3.5, num_inference_steps=20, lora_scale=0.8):
    # This function will be called by gr.load(), so we don't need to implement it here
    pass

examples = [
    ["d3xt3r as a camp counselor in the woods."],
    ["d3xt3r dressed as batman"],
    ["d3xt3r in a suit and tie"],
]

demo = gr.load("models/GenAIJake/d3xt3r")

# Customize the interface
with gr.Blocks() as new_demo:
    gr.Markdown("# D3XT3R Dachshund Image Generator")
    
    with gr.Row():
        with gr.Column():
            for component in demo.input_components:
                component.render()
            
            # Add LoRA scale slider
            lora_scale = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.05, label="LoRA Influence")
            
            generate_btn = gr.Button("Generate")
        
        with gr.Column():
            for component in demo.output_components:
                component.render()
    
    gr.Examples(examples, inputs=demo.input_components)
    
    # Update the generate button to include the lora_scale parameter
    generate_btn.click(
        fn=generate_image,
        inputs=demo.input_components + [lora_scale],
        outputs=demo.output_components
    )

# Add GPUZero functionality
new_demo.queue()
new_demo.launch()