Spaces:
Runtime error
Runtime error
File size: 2,237 Bytes
a24294c 04b10bd a485d72 a24294c 04b10bd a24294c 04b10bd 4b2f0f1 04b10bd a485d72 5d657f5 04b10bd a24294c 5d657f5 a24294c 5d657f5 a24294c 5d657f5 a24294c 5d657f5 a485d72 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import gradio as gr
import torch
from diffusers import DiffusionPipeline # Note: Change `FluxPipeline` to `DiffusionPipeline` if `FluxPipeline` is not correct
from PIL import Image
# Function to determine the device and handle model loading
def setup_pipeline():
# Check for CUDA availability
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the diffusion model
try:
pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
if device == "cpu":
# If using CPU, ensure model is offloaded to avoid GPU-specific features
pipeline.enable_model_cpu_offload()
else:
# Move model to GPU
pipeline.to(device)
except Exception as e:
print(f"Error loading model: {e}")
raise e
return pipeline, device
pipeline, device = setup_pipeline()
def generate_image(prompt, guidance_scale=7.5, num_inference_steps=50):
# Generate an image based on the prompt
with torch.no_grad():
try:
images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
except Exception as e:
print(f"Error generating image: {e}")
raise e
# Assuming pipeline returns a list of images, just take the first one
img = images[0]
# Convert PIL image to format suitable for Gradio
return img
# Set up Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Text to Image Generation")
with gr.Row():
prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
guidance_scale = gr.Slider(minimum=1, maximum=15, step=0.1, value=7.5, label="Guidance Scale")
num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=50, label="Number of Inference Steps")
with gr.Row():
generate_button = gr.Button("Generate Image")
result = gr.Image(label="Generated Image")
# Connect the function to the button
generate_button.click(
fn=generate_image,
inputs=[prompt, guidance_scale, num_inference_steps],
outputs=result
)
# Launch the app
demo.launch()
|