nightfury commited on
Commit
8549305
·
1 Parent(s): 70cff9c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -21
app.py CHANGED
@@ -1,42 +1,57 @@
1
  import gradio as gr
2
- import torch
3
- import streamlit as st
 
4
  from PIL import Image
5
  import numpy as np
6
  from io import BytesIO
 
 
 
7
  from diffusers import StableDiffusionImg2ImgPipeline
8
 
 
 
 
 
9
  device="cuda"
10
 
11
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=st.secrets['USER_TOKEN'])
12
- pipe.to(device)
 
 
 
 
 
 
13
 
14
  def resize(value,img):
 
15
  img = Image.open(img)
 
 
 
16
  img = img.resize((value,value), Image.Resampling.LANCZOS)
17
  return img
18
 
19
- def infer(source_img, prompt, guide, steps, seed, Strength):
20
- generator = torch.Generator('cuda').manual_seed(seed)
 
21
  source_image = resize(512, source_img)
22
  source_image.save('source.png')
23
- image_list = pipe([prompt], init_image=source_image, strength=Strength, guidance_scale=guide, num_inference_steps=steps)
24
  images = []
25
  safe_image = Image.open(r"unsafe.png")
26
- for i, image in enumerate(image_list["sample"]):
27
- if(image_list["nsfw_content_detected"][i]):
28
  images.append(safe_image)
29
  else:
30
  images.append(image)
31
- return image
32
-
33
- gr.Interface(fn=infer, inputs=[gr.Image(source="upload", type="filepath", label="Raw Image"), gr.Textbox(label = 'Prompt Input Text'),
34
- gr.Slider(2, 15, value = 7, label = 'Guidence Scale'),
35
- gr.Slider(10, 50, value = 25, step = 1, label = 'Number of Iterations'),
36
- gr.Slider(
37
- label = "Seed",
38
- minimum = 0,
39
- maximum = 2147483647,
40
- step = 1,
41
- randomize = True), gr.Slider(label='Strength', minimum = 0, maximum = 1, step = .05, value = .5)
42
- ], outputs='image').queue(max_size=10).launch(enable_queue=True)
 
1
  import gradio as gr
2
+ #import torch
3
+ #from torch import autocast // only for GPU
4
+
5
  from PIL import Image
6
  import numpy as np
7
  from io import BytesIO
8
+ import os
9
+ MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')
10
+
11
  from diffusers import StableDiffusionImg2ImgPipeline
12
 
13
+ print("hello !")
14
+
15
+ YOUR_TOKEN=MY_SECRET_TOKEN
16
+
17
  device="cuda"
18
 
19
+ #prompt_pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
20
+ #prompt_pipe.to(device)
21
+
22
+ img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
23
+ img_pipe.to(device)
24
+
25
+ source_img = gr.Image(source="upload", type="filepath", label="init_img | 512*512 px")
26
+ gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
27
 
28
  def resize(value,img):
29
+ #baseheight = value
30
  img = Image.open(img)
31
+ #hpercent = (baseheight/float(img.size[1]))
32
+ #wsize = int((float(img.size[0])*float(hpercent)))
33
+ #img = img.resize((wsize,baseheight), Image.Resampling.LANCZOS)
34
  img = img.resize((value,value), Image.Resampling.LANCZOS)
35
  return img
36
 
37
+
38
+ def infer(prompt, source_img):
39
+
40
  source_image = resize(512, source_img)
41
  source_image.save('source.png')
42
+ images_list = img_pipe([prompt] * 2, init_image=source_image, strength=0.75)
43
  images = []
44
  safe_image = Image.open(r"unsafe.png")
45
+ for i, image in enumerate(images_list["sample"]):
46
+ if(images_list["nsfw_content_detected"][i]):
47
  images.append(safe_image)
48
  else:
49
  images.append(image)
50
+ return images
51
+
52
+ print("Great ! Everything is working fine !")
53
+
54
+ title="Img2Img Stable Diffusion"
55
+ description="Img2Img Stable Diffusion example using CPU and HF token. <br />Warning: Slow process... ~5/10 min inference time. <b>NSFW filter enabled.</b>"
56
+
57
+ gr.Interface(fn=infer, inputs=["text", source_img], outputs=gallery,title=title,description=description).queue(max_size=100).launch(enable_queue=True)