gnosticdev commited on
Commit
65731df
·
verified ·
1 Parent(s): 3a5d85c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +83 -0
app.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import torch
4
+ from diffusers import StableDiffusionInpaintPipeline, StableDiffusionUpscalePipeline
5
+ import numpy as np
6
+
7
+ def process_image(image, prompt, mode, scale_factor=2):
8
+ if mode == "upscale":
9
+ # Upscale pipeline
10
+ pipeline = StableDiffusionUpscalePipeline.from_pretrained(
11
+ "stabilityai/stable-diffusion-x4-upscaler"
12
+ )
13
+ pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
14
+
15
+ # Process image
16
+ upscaled_image = pipeline(
17
+ prompt=prompt,
18
+ image=image,
19
+ noise_level=20,
20
+ num_inference_steps=20
21
+ ).images[0]
22
+
23
+ return upscaled_image
24
+
25
+ elif mode == "inpaint":
26
+ # Inpainting pipeline
27
+ pipeline = StableDiffusionInpaintPipeline.from_pretrained(
28
+ "runwayml/stable-diffusion-inpainting"
29
+ )
30
+ pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
31
+
32
+ # Create mask for extending the image
33
+ width, height = image.size
34
+ mask = Image.new('RGB', (width, height), 'white')
35
+
36
+ # Process image
37
+ result = pipeline(
38
+ prompt=prompt,
39
+ image=image,
40
+ mask_image=mask,
41
+ num_inference_steps=20
42
+ ).images[0]
43
+
44
+ return result
45
+
46
+ # Gradio Interface
47
+ def create_interface():
48
+ with gr.Blocks(title="AI Image Enhancement") as interface:
49
+ gr.Markdown("# AI Image Enhancement Studio")
50
+ gr.Markdown("Enhance, upscale, and recreate images using AI")
51
+
52
+ with gr.Row():
53
+ with gr.Column():
54
+ input_image = gr.Image(type="pil", label="Upload Image")
55
+ prompt = gr.Textbox(label="Prompt", placeholder="Describe the desired enhancement...")
56
+ mode = gr.Radio(
57
+ choices=["upscale", "inpaint"],
58
+ label="Processing Mode",
59
+ value="upscale"
60
+ )
61
+ scale_factor = gr.Slider(
62
+ minimum=2,
63
+ maximum=8,
64
+ step=2,
65
+ label="Upscale Factor",
66
+ value=2
67
+ )
68
+ process_btn = gr.Button("Process Image")
69
+
70
+ with gr.Column():
71
+ output_image = gr.Image(type="pil", label="Enhanced Result")
72
+
73
+ process_btn.click(
74
+ fn=process_image,
75
+ inputs=[input_image, prompt, mode, scale_factor],
76
+ outputs=output_image
77
+ )
78
+
79
+ return interface
80
+
81
+ if __name__ == "__main__":
82
+ interface = create_interface()
83
+ interface.launch()