feat: refactor generate_image under module flux

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import requests
 
+import flux
+
 def greet(name: str) -> str:
     """
     Greets the user with a personalized hello message.
@@ -69,21 +71,6 @@ def calculate(n1, op, n2):
     if op == "/" and n2 != 0: return str(n1 / n2)
     return "Error"
     
-def generate_image(prompt: str):
-    """
-    Generates an image based on a text prompt using the Stable Diffusion model.
-
-    Args:
-        prompt (str): The text prompt to generate the image from.
-
-    Returns:
-        str: The URL of the generated image.
-    """
-
-    # Call the Stable Diffusion API or model here
-    # For demonstration, we'll return a placeholder image URL
-    return "https://avatars.githubusercontent.com/u/75182?v=4"
-
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -133,11 +120,17 @@ with gr.Blocks() as demo:
     with tab_image:
         img_prompt = gr.Textbox(
             label="Image Prompt",
-            placeholder="Enter a prompt for the image generation"
+            placeholder="Enter a prompt for the image generation",
+            value="A portrait of a handsome software developer"
         )
         img_btn = gr.Button("Generate Image")
         img_output = gr.Image(label="Sample Image")
-        img_btn.click(fn=generate_image, inputs=img_prompt, outputs=img_output, api_name="generate_image")
+        img_btn.click(
+            fn=flux.generate_image,
+            # fn=flux.generate_image2, # DEBUG
+            inputs=img_prompt,
+            outputs=img_output,
+            api_name="generate_image")
 
 
 demo.launch(mcp_server=True, share=True)

flux.py

@@ -8,6 +8,7 @@ import time
 from io import BytesIO
 from pathlib import Path
 import modal
+import modal.running_app
 
 # We’ll make use of the full CUDA toolkit in this example, so we’ll build our container image
 # off of the nvidia/cuda base.
@@ -217,6 +218,101 @@ def optimize(pipe, compile=True):
 
     return pipe
 
+
+def generate_image(prompt: str):
+    """
+    Generates an image based on a text prompt using the Flux model.
+
+    Args:
+        prompt (str): The text prompt to generate the image from.
+
+    Returns:
+        str: The URL of the generated image.
+    """
+
+    # For demonstration, we'll return a placeholder image URL
+
+    return "https://avatars.githubusercontent.com/u/75182?v=4"
+
+
+app_running = False
+
+def make_sure_app_running():
+    """
+    Ensure the Modal app is running.
+
+    NOTE: WORK-IN-PROGRESS
+    """
+    global app_running
+    # In practice, you would implement logic to check if the Modal app is running.
+    # For example, you could use `modal app list` to check if the app is running.
+    if not app_running:
+        print("⚠️ The Modal app is not running. Should start the app first.")
+        print("Starting Modal app...")
+        # Start the Modal app
+        # This will run the app in the background and allow it to handle requests.
+        # Note: This is a blocking call, so it will not return until the app is stopped.
+        # If you want to run it in the background, you can use `modal run --detach`.
+        # For example:
+        # modal run flux.py --prompt "a beautiful landscape with mountains and a river" --twice --compile
+        app.run(detach=False, interactive=True)
+        print("✅ Modal app started successfully.")
+        app_running = True
+        print("You can now call the generate_image function to generate images.")
+        print("Example usage:")
+        print("generate_image(prompt='a beautiful landscape with mountains and a river', twice=True, compile=False)")
+    else:
+        print("✅ The Modal app is already running. You can call the generate_image function to generate images.")
+    pass
+
+
+def generate_image2(    
+    prompt: str = """
+    A portrait of a handsome software developer
+    """,
+    twice: bool = True,
+    compile: bool = False,
+):
+    """
+    Generates an image based on a text prompt using the Flux model
+    running on Modal serverless infrastructure.
+
+    Args:
+        prompt (str): The text prompt to generate the image from.
+        twice (bool): Whether to run the inference twice.
+        compile (bool): Whether to compile the model.
+
+    Returns:
+        str: The URL of the generated image.
+    """
+
+    print("DEBUG: generate_image2 called with parameters:")
+    print(f"  prompt: {prompt}")
+    print(f"  twice: {twice}")
+    print(f"  compile: {compile}")
+    print("DEBUG: Starting image generation...")
+
+    # Ensure the Modal app is running before making any calls
+    make_sure_app_running()
+
+    t0 = time.time()
+    image_bytes = Model(compile=compile).inference.remote(prompt)
+    print(f"🎨 first inference latency: {time.time() - t0:.2f} seconds")
+
+    if twice:
+        t0 = time.time()
+        image_bytes = Model(compile=compile).inference.remote(prompt)
+        print(f"🎨 second inference latency: {time.time() - t0:.2f} seconds")
+
+    output_path = Path("/tmp") / "flux" / "output.jpg"
+    output_path.parent.mkdir(exist_ok=True, parents=True)
+    print(f"🎨 saving output to {output_path}")
+    output_path.write_bytes(image_bytes)
+
+    print(f"✅ Image generated and saved to {output_path}")
+    return output_path
+
+
 # To run this script, use the command:
 # modal run flux.py --prompt "a beautiful landscape with mountains and a river" --twice --compile