Update app.py

app.py

@@ -5,13 +5,17 @@ import spaces
 from diffusers import FluxControlPipeline, FluxTransformer2DModel
 
 ####################################
-#   Load the model(s) on CPU       #
+#   Load the model(s) on GPU       #
 ####################################
-path = "sayakpaul/FLUX.1-dev-edit-v0" 
+path = "sayakpaul/FLUX.1-dev-edit-v0"
 edit_transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16)
 pipeline = FluxControlPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", transformer=edit_transformer, torch_dtype=torch.bfloat16
 ).to("cuda")
+pipeline.load_lora_weights(
+    hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd"
+)
+pipeline.set_adapters(["hyper-sd"], adapter_weights=[0.125])
 
 #####################################
 #  The function for our Gradio app  #
@@ -24,14 +28,14 @@ def generate(prompt, input_image):
     """
     # Perform inference
     output_image = pipeline(
-        control_image=input_image,
+        control_image=image,
         prompt=prompt,
-        guidance_scale=30.0,
-        num_inference_steps=50,
+        guidance_scale=30.,
+        num_inference_steps=8,
         max_sequence_length=512,
-        height=input_image.height,
-        width=input_image.width,
-        generator=torch.manual_seed(0),
+        height=image.height,
+        width=image.width,
+        generator=torch.manual_seed(0)
     ).images[0]
 
     return output_image
@@ -41,7 +45,7 @@ def launch_app():
     with gr.Blocks() as demo:
         gr.Markdown(
             """
-            # Flux Control Editing
+            # Flux Control Editing 🖌️
 
             This demo uses the [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) 
             pipeline with an edit transformer from [Sayak Paul](https://huggingface.co/sayakpaul).