AP123 commited on
Commit
a29e3ba
β€’
1 Parent(s): 70d6a02

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -22
app.py CHANGED
@@ -3,35 +3,28 @@ import os
3
  import gradio as gr
4
  from PIL import Image
5
  from diffusers import (
6
- StableDiffusionPipeline,
7
  StableDiffusionControlNetImg2ImgPipeline,
8
  ControlNetModel,
9
- DDIMScheduler,
10
- DPMSolverMultistepScheduler,
11
- DEISMultistepScheduler,
12
- HeunDiscreteScheduler,
13
- EulerDiscreteScheduler,
14
  )
15
 
16
- # Load controlnet model in float16 precision
17
- controlnet = ControlNetModel.from_pretrained(
18
- "monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16
19
- )
20
-
21
- # Load the pipeline in float16 precision
22
- pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
23
  "SG161222/Realistic_Vision_V2.0",
24
  controlnet=controlnet,
25
  safety_checker=None,
26
  torch_dtype=torch.float16,
27
  ).to("cuda")
28
- pipe.enable_xformers_memory_efficient_attention()
29
 
 
30
  SAMPLER_MAP = {
31
  "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
32
  "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
33
  }
34
 
 
35
  def inference(
36
  control_image: Image.Image,
37
  prompt: str,
@@ -45,16 +38,15 @@ def inference(
45
  if prompt is None or prompt == "":
46
  raise gr.Error("Prompt is required")
47
 
48
- # Generate init_image using the "Realistic Vision V2.0" model
49
- init_image = pipe(prompt, height=512, width=512).images[0]
50
- print("Init Image:", init_image)
51
- assert init_image is not None, "init_image is None!"
52
-
53
  control_image = control_image.resize((512, 512))
54
- pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)
55
  generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()
56
-
57
- out = pipe(
58
  prompt=prompt,
59
  negative_prompt=negative_prompt,
60
  image=init_image,
 
3
  import gradio as gr
4
  from PIL import Image
5
  from diffusers import (
6
+ DiffusionPipeline,
7
  StableDiffusionControlNetImg2ImgPipeline,
8
  ControlNetModel,
 
 
 
 
 
9
  )
10
 
11
+ # Initialize both pipelines
12
+ init_pipe = DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V2.0", torch_dtype=torch.float16).to("cuda")
13
+ controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)
14
+ main_pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
 
 
 
15
  "SG161222/Realistic_Vision_V2.0",
16
  controlnet=controlnet,
17
  safety_checker=None,
18
  torch_dtype=torch.float16,
19
  ).to("cuda")
 
20
 
21
+ # Sampler map
22
  SAMPLER_MAP = {
23
  "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
24
  "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
25
  }
26
 
27
+ # Inference function
28
  def inference(
29
  control_image: Image.Image,
30
  prompt: str,
 
38
  if prompt is None or prompt == "":
39
  raise gr.Error("Prompt is required")
40
 
41
+ # Generate the initial image
42
+ init_image = init_pipe(prompt).images[0]
43
+
44
+ # Rest of your existing code
 
45
  control_image = control_image.resize((512, 512))
46
+ main_pipe.scheduler = SAMPLER_MAP[sampler](main_pipe.scheduler.config)
47
  generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()
48
+
49
+ out = main_pipe(
50
  prompt=prompt,
51
  negative_prompt=negative_prompt,
52
  image=init_image,