Create app.py

app.py

@@ -0,0 +1,54 @@
+import torch
+from PIL import Image
+import gradio as gr
+from diffusers import StableDiffusionImg2ImgPipeline
+
+# Load the Stable Diffusion img2img pipeline
+device = "cuda"
+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()