Create app.py

app.py

@@ -0,0 +1,45 @@
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
+
+# Function to load the model and list layers
+def list_flux_layers():
+    try:
+        # Load the FLUX.1-dev model
+        # Using torch.bfloat16 to reduce memory usage, offloading to CPU if needed
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.bfloat16
+        )
+        pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM
+
+        # Access the transformer (main component of Flux) and get layer names
+        model = pipe.transformer  # Flux's core transformer model
+        layer_names = []
+
+        # Iterate through all named modules to get layer names
+        for name, module in model.named_modules():
+            layer_names.append(name)
+
+        # Format the output as a numbered list
+        output = "\n".join([f"{i+1}. {name}" for i, name in enumerate(layer_names)])
+        return f"Layers of FLUX.1-dev:\n\n{output}"
+    except Exception as e:
+        return f"Error loading model or listing layers: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="FLUX.1-dev Layer Lister") as demo:
+    gr.Markdown("# FLUX.1-dev Layer Lister")
+    gr.Markdown("Click the button below to list all layers in the black-forest-labs/FLUX.1-dev model.")
+    
+    # Button to trigger the layer listing
+    btn = gr.Button("List Layers")
+    
+    # Output area for the layer names
+    output = gr.Textbox(label="Model Layers", lines=20, placeholder="Layer names will appear here...")
+    
+    # Connect the button to the function
+    btn.click(fn=list_flux_layers, inputs=None, outputs=output)
+
+# Launch the app
+demo.launch()