Spaces:
Running
Running
File size: 2,154 Bytes
d5b83ed 62831dd d5b83ed 36beca2 d5b83ed 62831dd d5b83ed |
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 |
import torch
from PIL import Image
import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
# Load the Stable Diffusion img2img pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
).to(device)
#pipe.enable_attention_slicing()
# Predefined prompts
PROMPTS = {
"Photorealistic": "A heavily corroded metal pipeline stretching across the ocean floor, covered in algae and barnacles, deep underwater with ambient blue lighting and floating particles, photorealistic",
"Cinematic": "An old rusted pipeline submerged in the sea, encrusted with marine growth and decay, surrounded by dark water and shafts of light from the surface, cinematic, moody atmosphere"
}
# Inference function
def generate_image(init_image, prompt_choice, strength, guidance_scale):
# Resize and convert the input image
init_image = init_image.convert("RGB").resize((768, 512))
# Get the selected prompt
prompt = PROMPTS[prompt_choice]
# Run the pipeline
result = pipe(
prompt=prompt,
image=init_image,
strength=strength,
guidance_scale=guidance_scale
).images[0]
return result
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("Corroded Pipeline Generator - Underwater Img2Img")
with gr.Row():
with gr.Column():
init_image = gr.Image(label="Upload Initial Image", type="pil")
prompt_choice = gr.Radio(choices=list(PROMPTS.keys()), label="Select Prompt", value="Photorealistic")
strength = gr.Slider(minimum=0.2, maximum=1.0, value=0.75, step=0.05, label="Transformation Strength")
guidance_scale = gr.Slider(minimum=1, maximum=15, value=7.5, step=0.5, label="Prompt Guidance Scale")
generate_btn = gr.Button("Generate")
with gr.Column():
output_image = gr.Image(label="Generated Image")
generate_btn.click(fn=generate_image, inputs=[init_image, prompt_choice, strength, guidance_scale], outputs=output_image)
# Launch the app
demo.launch()
|