change layout to lost one model

app.py

@@ -11,8 +11,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "black-forest-labs/FLUX.1-schnell"  # Replace to the model you would like to use
 torch_dtype = torch.bfloat16
 
-pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# pipe = pipe.to(device)
 
 # load pruned model
 pruned_pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
@@ -29,14 +29,12 @@ MAX_IMAGE_SIZE = 1024
 
 @spaces.GPU
 def generate_images(prompt, seed, steps):
-    pipe.to("cuda")
-    pruned_pipe.to("cuda")
     # Run the model and return images directly
-    g_cpu = torch.Generator("cuda").manual_seed(seed)
-    original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
+    # g_cpu = torch.Generator("cuda").manual_seed(seed)
+    # original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
     g_cpu = torch.Generator("cuda").manual_seed(seed)
     ecodiff_image = pruned_pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
-    return original_image, ecodiff_image
+    return ecodiff_image
 
 
 examples = [
@@ -56,6 +54,8 @@ css = """
 
 header = """
 # 🌱 EcoDiff Pruned FLUX-Schnell (20% Pruning Ratio)
+
+We are not able to host two FLUX models in the same space, so we only show the pruned model here.
 """
 
 header_2 = """
@@ -95,7 +95,7 @@ with gr.Blocks(css=css) as demo:
         inputs=[prompt],
     )
     with gr.Row():
-        original_output = gr.Image(label="Original Output")
+        # original_output = gr.Image(label="Original Output")
         ecodiff_output = gr.Image(label="EcoDiff Output")
     gr.on(
         triggers=[generate_btn.click, prompt.submit],
@@ -105,7 +105,7 @@ with gr.Blocks(css=css) as demo:
             seed,
             steps,
         ],
-        outputs=[original_output, ecodiff_output],
+        outputs=[ecodiff_output],
     )
 
 if __name__ == "__main__":