import gradio as gr
from diffusers import DiffusionPipeline
import torch
from PIL import Image

# Load the pipeline
prj_path = "jkcg/furniture-chair"
model = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(
     model,
     torch_dtype=torch.float32,  # Use float32 for CPU
)
pipe.to("cpu")  # Ensure the pipeline runs on CPU
pipe.load_lora_weights(prj_path, weight_name="pytorch_lora_weights.safetensors")

def generate_image(prompt, seed):
    generator = torch.Generator("cpu").manual_seed(seed)
    image = pipe(prompt=prompt, generator=generator).images[0]
    return image

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="photo of a furnichair-texx in an empty room"),
        gr.Slider(label="Seed", minimum=0, maximum=10000, step=1, value=42)
    ],
    outputs=gr.Image(label="Generated Image")
)

# Launch the interface
interface.launch(share=True)