import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch

def load_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    pipeline = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32  # Adjusted for device compatibility
    ).to(device)
    
    # Ensure this line points to the correct path of your weights
    pipeline.load_lora_weights("./", weight_name="EnvyFloorplansXL01.safetensors")
    
    return pipeline

pipeline = load_model()

def generate_image(prompt, negative_prompt):
    # Generate an image using the model pipeline
    image = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        scheduler="DPM++ 2M Karras Sharp v1",
        generator=torch.manual_seed(3145577831),
        num_inference_steps=60,
        guidance_scale=9.5
    ).images[0]
    return image

# Define your Gradio interface
# Define your Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here...", value="(poor quality:1.2) (worst quality, low quality:1.4),word,cropped,username,watermark,signature,blurry,soft,soft line,curved line,sketch,ugly,logo,pixelated,lowres,ceiling light, monochrome, negativeXL_D, color pencil, soft line,worst quality, blurry,sketch,, text,, sketch, cartoon, drawing")  # Corrected here
    ],
    outputs="image",
    title="Text-to-Image Generation with Stable Diffusion",
    description="Generate images based on your text inputs."
)


if __name__ == "__main__":
    interface.launch(debug=True)