import torch from diffusers import DiffusionPipeline import gradio as gr # Load the Flux model with ControlNet upscaling pipe = DiffusionPipeline.from_pretrained("jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.float16) pipe.to("cuda") # Send the model to GPU if available # Function to generate an image based on user prompt def generate_image(prompt): # Generate an image using the Flux model with torch.inference_mode(): image = pipe(prompt).images[0] return image # Set up the Gradio interface with gr.Blocks() as demo: gr.Markdown("# Image Generation using Flux (jasperai/Flux.1-dev-Controlnet-Upscaler)") # Input for user to provide a text prompt prompt_input = gr.Textbox(label="Enter Text Prompt", placeholder="e.g. 'Astronaut in a jungle, cold color palette, muted colors, detailed, 8k'", value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k") # Output to display the generated image output_image = gr.Image(label="Generated Image") # Button to trigger image generation generate_button = gr.Button("Generate Image") # Link button click to image generation function generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image) # Launch the Gradio app demo.launch()