Create main.py
Browse files
main.py
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import gradio as gr
|
3 |
+
from PIL import Image
|
4 |
+
from diffusers import StableDiffusionPipeline
|
5 |
+
|
6 |
+
# Use a pipeline as a high-level helper
|
7 |
+
from transformers import pipeline
|
8 |
+
|
9 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
10 |
+
|
11 |
+
caption_image = pipeline("image-to-text",
|
12 |
+
model="Salesforce/blip-image-captioning-large", device=device)
|
13 |
+
|
14 |
+
|
15 |
+
def image_generation(prompt):
|
16 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
17 |
+
pipeline = StableDiffusionPipeline.from_pretrained(
|
18 |
+
"stabilityai/stable-diffusion-3-medium",
|
19 |
+
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
20 |
+
)
|
21 |
+
#pipeline.to(device)
|
22 |
+
pipeline.enable_model_cpu_offload()
|
23 |
+
|
24 |
+
image = pipeline(
|
25 |
+
prompt=prompt + " 8K, Ultra HD",
|
26 |
+
negative_prompt="blurred, ugly, watermark, low resolution, blurry, nude",
|
27 |
+
num_inference_steps=40,
|
28 |
+
height=1024,
|
29 |
+
width=1024,
|
30 |
+
guidance_scale=9.0
|
31 |
+
).images[0]
|
32 |
+
|
33 |
+
return image
|
34 |
+
|
35 |
+
def caption_my_image(pil_image):
|
36 |
+
semantics = caption_image(images=pil_image)[0]['generated_text']
|
37 |
+
images = image_generation(semantics)
|
38 |
+
return images
|
39 |
+
|
40 |
+
demo = gr.Interface(fn=caption_my_image,
|
41 |
+
inputs=[gr.Image(label="Select Image",type="pil")],
|
42 |
+
outputs=[gr.Image(label="New Image genrated using SD3",type="pil")],
|
43 |
+
title="PicTalker | ImageNarrator | SnapSpeech | SpeakScene",
|
44 |
+
description="🌟 Transform Ordinary Photos into Extraordinary Art!")
|
45 |
+
demo.launch()
|