debug: base model + output management (#2)

- debug: base model + output management (7f98d3f3499728559c1602b3c40e2e9c433682e7)


Co-authored-by: Mehdi Merabet <[email protected]>

app.py

@@ -1,14 +1,17 @@
 import gradio as gr
 from diffusers import DiffusionPipeline
 
-# Load the model
-pipe = DiffusionPipeline.from_pretrained("Yaquv/rick")
+# Load the base model
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model)
 
-# Define a prediction function to handle the output
+# Load the LoRA weights from your repository
+pipe.load_lora_weights("Yaquv/rick", weight_name="rick.safetensors")
+
+# Define a prediction function
 def predict(text):
-    output = pipe(text)  # Run the model with the text input
-    # Assuming the output is a tuple and the image is the first element
-    image = output[0]
+    output = pipe(text)
+    image = output.images[0]  # Extract the first image from the output
     return image
 
 # Create a Gradio interface