Anurag181011 commited on
Commit
49c57cc
·
1 Parent(s): 4e7c0a2
Files changed (1) hide show
  1. app.py +18 -13
app.py CHANGED
@@ -3,38 +3,43 @@ import torch
3
  from diffusers import StableDiffusionImg2ImgPipeline
4
  from PIL import Image
5
 
6
- # Check if GPU is available
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
  print(f"Using device: {device}")
 
9
 
10
- # Load the model correctly
 
 
 
 
 
 
 
11
  model_id = "nitrosocke/Ghibli-Diffusion"
12
  pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
13
- model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
14
  )
15
 
16
- # Enable optimization if using CUDA
17
  if device == "cuda":
18
  pipe.to(device)
19
- pipe.enable_model_cpu_offload() # Helps with memory efficiency
20
- else:
21
- pipe.to("cpu") # Ensure CPU compatibility
22
 
23
  def transform_image(input_image: Image.Image) -> Image.Image:
24
  input_image = input_image.resize((512, 512))
25
-
26
- # Move image to proper format for inference
27
  prompt = "ghibli style, cinematic lighting, hand-painted, anime aesthetics"
 
28
  output = pipe(
29
  prompt=prompt,
30
  image=input_image,
31
- strength=0.75,
32
- guidance_scale=7.5,
33
- num_inference_steps=30, # Reduced steps for faster generation
34
  )
35
  return output.images[0]
36
 
37
- # Build the Gradio UI
38
  demo = gr.Interface(
39
  fn=transform_image,
40
  inputs=gr.Image(type="pil", label="Upload your portrait/photo"),
 
3
  from diffusers import StableDiffusionImg2ImgPipeline
4
  from PIL import Image
5
 
6
+ # Check and print device status
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
  print(f"Using device: {device}")
9
+ print(f"CUDA available: {torch.cuda.is_available()}")
10
 
11
+ # Ensure torch is installed correctly
12
+ try:
13
+ torch.zeros(1).to(device)
14
+ print("Torch initialized successfully on", device)
15
+ except Exception as e:
16
+ print("Torch initialization error:", e)
17
+
18
+ # Load the Stable Diffusion model with optimizations
19
  model_id = "nitrosocke/Ghibli-Diffusion"
20
  pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
21
+ model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32, safety_checker=None
22
  )
23
 
 
24
  if device == "cuda":
25
  pipe.to(device)
26
+ pipe.enable_model_cpu_offload() # Efficient VRAM usage
27
+ pipe.enable_xformers_memory_efficient_attention() # Optimized attention for speed
 
28
 
29
  def transform_image(input_image: Image.Image) -> Image.Image:
30
  input_image = input_image.resize((512, 512))
 
 
31
  prompt = "ghibli style, cinematic lighting, hand-painted, anime aesthetics"
32
+
33
  output = pipe(
34
  prompt=prompt,
35
  image=input_image,
36
+ strength=0.65, # Reduce strength to avoid excessive details
37
+ guidance_scale=5.0, # Lowered for faster inference
38
+ num_inference_steps=25, # Reduced steps for speed
39
  )
40
  return output.images[0]
41
 
42
+ # Gradio Interface
43
  demo = gr.Interface(
44
  fn=transform_image,
45
  inputs=gr.Image(type="pil", label="Upload your portrait/photo"),