Create app.py

app.py

@@ -0,0 +1,61 @@
+import spaces
+from diffusers import AuraFlowPipeline
+import torch
+import gradio as gr
+
+def initialize_auraflow_pipeline():
+    """Initialize and return the AuraFlowPipeline."""
+    pipeline = AuraFlowPipeline.from_pretrained(
+        "fal/AuraFlow-v0.3",
+        torch_dtype=torch.float16,
+        variant="fp16",
+    ).to("cuda")
+    return pipeline
+
+def generate_image(pipeline, prompt, width, height, num_inference_steps, seed, guidance_scale):
+    """Generate an image using the AuraFlowPipeline."""
+    generator = torch.Generator().manual_seed(seed)
+
+    image = pipeline(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=guidance_scale,
+    ).images[0]
+
+    return image
+
+# Initialize the pipeline once
+auraflow_pipeline = initialize_auraflow_pipeline()
+
+# Gradio interface
+def gradio_generate_image(prompt, width, height, num_inference_steps, seed, guidance_scale):
+    return generate_image(auraflow_pipeline, prompt, width, height, num_inference_steps, seed, guidance_scale)
+
+# Create Gradio Blocks
+with gr.Blocks() as demo:
+    gr.Markdown("# AuraFlow Image Generation")
+    
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your image prompt here...")
+            width_input = gr.Slider(minimum=512, maximum=2048, step=64, value=1536, label="Width")
+            height_input = gr.Slider(minimum=512, maximum=2048, step=64, value=768, label="Height")
+            steps_input = gr.Slider(minimum=10, maximum=100, step=1, value=50, label="Number of Inference Steps")
+            seed_input = gr.Number(label="Seed", value=1)
+            guidance_input = gr.Slider(minimum=1, maximum=10, step=0.1, value=3.5, label="Guidance Scale")
+            generate_btn = gr.Button("Generate Image")
+        
+        with gr.Column():
+            image_output = gr.Image(label="Generated Image")
+    
+    generate_btn.click(
+        fn=gradio_generate_image,
+        inputs=[prompt_input, width_input, height_input, steps_input, seed_input, guidance_input],
+        outputs=image_output
+    )
+
+# Launch the Gradio interface
+demo.launch()