Spaces:
Build error
Build error
File size: 5,166 Bytes
002bf33 49ed47b d4c4d41 002bf33 b928aa3 002bf33 49ed47b 002bf33 49ed47b 002bf33 eca5a0d 002bf33 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | # Import necessary libraries
import torch
import gradio as gr
import webbrowser
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image
# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
print("⚠️ Warning: Running on CPU, performance may be slow.")
# Load Text-to-Image model
print("🔄 Loading Stable Diffusion txt2img model...")
pipe_txt2img = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)
print("✅ Text-to-Image model loaded!")
# Load Image-to-Image model
print("🔄 Loading Stable Diffusion img2img model...")
pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)
print("✅ Image-to-Image model loaded!"
)
# Function to generate images from text
def generate_txt2img(prompt, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
generator = torch.manual_seed(seed) if seed != -1 else None
image = pipe_txt2img(
prompt, num_inference_steps=steps, guidance_scale=guidance, width=width, height=height,
generator=generator
).images[0]
output_path = f"generated_image.{save_format}" # Save image in the requested format
image.save(output_path, format=save_format.upper()) # Save in the selected format
print(f"Image saved to {output_path}")
return output_path
# Function to transform images using img2img
def generate_img2img(prompt, image, strength=0.5, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
generator = torch.manual_seed(seed) if seed != -1 else None
image = pipe_img2img(
prompt, image=image, strength=strength, num_inference_steps=steps, guidance_scale=guidance,
width=width, height=height, generator=generator
).images[0]
output_path = f"modified_image.{save_format}" # Save image in the requested format
image.save(output_path, format=save_format.upper()) # Save in the selected format
print(f"Image saved to {output_path}")
return output_path
# Define Gradio UI
def create_ui():
with gr.Blocks(title="DiffuGen: AI Image Generation") as demo:
gr.Markdown("# 🌟 DiffuGen - AI Image Generator")
# Text-to-Image Tab
with gr.Tab("📷 Text to Image"):
with gr.Row():
prompt = gr.Textbox(label="Enter a text prompt")
with gr.Row():
steps = gr.Slider(10, 100, value=50, step=10, label="Steps")
guidance = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
with gr.Row():
width = gr.Slider(256, 1024, value=512, step=64, label="Width")
height = gr.Slider(256, 1024, value=512, step=64, label="Height")
seed = gr.Number(value=-1, label="Seed (-1 for random)")
with gr.Row():
save_format = gr.Dropdown(
choices=["png", "jpg"], value="png", label="Select Image Format"
)
generate_btn = gr.Button("🚀 Generate Image")
output_image = gr.Image(label="Generated Image", type="pil")
generate_btn.click(
generate_txt2img,
inputs=[prompt, steps, guidance, width, height, seed, save_format],
outputs=output_image
)
# Image-to-Image Tab
with gr.Tab("🖼️ Image to Image"):
with gr.Row():
prompt_img2img = gr.Textbox(label="Enter a prompt")
with gr.Row():
input_img = gr.Image(label="Upload Image", type="pil")
with gr.Row():
strength = gr.Slider(0.1, 1.0, value=0.5, label="Denoising Strength")
steps_img2img = gr.Slider(10, 100, value=50, label="Steps")
guidance_img2img = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
with gr.Row():
width_img2img = gr.Slider(256, 1024, value=512, step=64, label="Width")
height_img2img = gr.Slider(256, 1024, value=512, step=64, label="Height")
seed_img2img = gr.Number(value=-1, label="Seed (-1 for random)")
with gr.Row():
save_format_img2img = gr.Dropdown(
choices=["png", "jpg"], value="png", label="Select Image Format"
)
generate_img_btn = gr.Button("🔄 Transform Image")
output_img2img = gr.Image(label="Modified Image", type="pil")
generate_img_btn.click(
generate_img2img,
inputs=[prompt_img2img, input_img, strength, steps_img2img, guidance_img2img, width_img2img, height_img2img, seed_img2img, save_format_img2img],
outputs=output_img2img
)
return demo
# Launch Gradio WebUI
web_ui = create_ui()
url = web_ui.launch(share=True)
# Automatically open the WebUI in a new browser tab
webbrowser.open(url)
|