Update app.py

app.py

@@ -1,25 +1,31 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from PIL import Image
-import numpy as np
+from diffusers import StableDiffusionPipeline
+from safetensors.torch import load_file
+import torch
+
+def load_pipeline(base_model_path, lora_path):
+    # Load base SD-v1.5 pipeline
+    pipeline = StableDiffusionPipeline.from_pretrained(base_model_path)
+
+    # Load LoRA weights
+    lora_state_dict = load_file(lora_path)
+    pipeline.unet.load_attn_procs(lora_state_dict)
 
-# Load the model
-model_name = "maria26/Floor_Plan_LoRA"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-model.eval()
+    return pipeline
 
-def predict(input_data):
-    # Tokenize input
-    inputs = tokenizer(input_data, return_tensors="pt")
+# Paths to the base model and LoRA weights
+base_model_path = "path/to/sd-v1-5"
+lora_path = "path/to/Floor_Plan_LoRA.safetensors"
 
-    # Generate output
-    outputs = model.generate(inputs["input_ids"], max_length=50)
+# Load the pipeline
+pipeline = load_pipeline(base_model_path, lora_path)
 
-    # Postprocess the output (replace with actual model output)
-    image_array = np.random.rand(128, 128, 3) * 255  # Placeholder for image
-    return Image.fromarray(image_array.astype("uint8"))
+def predict(prompt):
+    # Generate an image based on the prompt
+    result = pipeline(prompt).images[0]
+    return result
+
+# Create Gradio Interface
+import gradio as gr
 
-# Define Gradio interface
 interface = gr.Interface(fn=predict, inputs="text", outputs="image")
-interface.launch()
+interface.launch()