Update app.py

app.py

@@ -7,6 +7,9 @@ import spaces
 from PIL import Image
 from fastapi import FastAPI
 from fastapi.responses import FileResponse
+import transformers 
+
+transformers.utils.move_cache()
 
 # Initialize the base model
 dtype = torch.bfloat16
@@ -20,7 +23,7 @@ MAX_SEED = 2**32-1
 app = FastAPI()
 
 @spaces.GPU()
-def generate_image(prompt, width, height, lora_path, trigger_word, hash_value):
+def generate_image(prompt, width, height, lora_path, trigger_word, hash_value, steps):
     # Load LoRA weights
     pipe.load_lora_weights(lora_path)
     
@@ -34,7 +37,7 @@ def generate_image(prompt, width, height, lora_path, trigger_word, hash_value):
     # Generate image
     image = pipe(
         prompt=full_prompt,
-        num_inference_steps=28,
+        num_inference_steps=steps,
         guidance_scale=3.5,
         width=width,
         height=height,
@@ -50,8 +53,8 @@ def generate_image(prompt, width, height, lora_path, trigger_word, hash_value):
     
     return image, image_path
 
-def run_lora(prompt, width, height, lora_path, trigger_word, hash_value):
-    image, image_path = generate_image(prompt, width, height, lora_path, trigger_word, hash_value)
+def run_lora(prompt, width, height, lora_path, trigger_word, hash_value, steps):
+    image, image_path = generate_image(prompt, width, height, lora_path, trigger_word, hash_value, steps)
     return image, image_path
 
 # Set up the Gradio interface
@@ -62,8 +65,9 @@ with gr.Blocks() as gradio_app:
         prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here")
     
     with gr.Row():
-        width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=512)
-        height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=512)
+        width = gr.Slider(label="Width", minimum=128, maximum=1024, step=64, value=512)
+        height = gr.Slider(label="Height", minimum=128, maximum=1024, step=64, value=512)
+        steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=28)
     
     with gr.Row():
         lora_path = gr.Textbox(label="LoRA Path", value="SebastianBodza/Flux_Aquarell_Watercolor_v2")
@@ -77,7 +81,7 @@ with gr.Blocks() as gradio_app:
     
     generate_button.click(
         fn=run_lora,
-        inputs=[prompt, width, height, lora_path, trigger_word, hash_value],
+        inputs=[prompt, width, height, lora_path, trigger_word, hash_value, steps],
         outputs=[output_image, output_path]
     )