gokaygokay commited on
Commit
461d979
1 Parent(s): bc62600

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -3
app.py CHANGED
@@ -8,10 +8,12 @@ from diffusers import EulerAncestralDiscreteScheduler, DPMSolverMultistepSchedul
8
  from diffusers.models.attention_processor import AttnProcessor2_0
9
  import gradio as gr
10
  from PIL import Image
 
11
  from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
12
  import requests
13
  from RealESRGAN import RealESRGAN
14
 
 
15
  import subprocess
16
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
17
 
@@ -113,14 +115,19 @@ def enhance_prompt(input_prompt, model_choice):
113
 
114
  return enhanced_text
115
 
116
- # Upscale function
117
  def upscale_image(image, scale):
 
 
 
118
  if scale == 2:
119
- return realesrgan_x2.predict(image)
120
  elif scale == 4:
121
- return realesrgan_x4.predict(image)
122
  else:
123
  return image
 
 
 
124
 
125
  @spaces.GPU(duration=120)
126
  def generate_image(additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps,
 
8
  from diffusers.models.attention_processor import AttnProcessor2_0
9
  import gradio as gr
10
  from PIL import Image
11
+ import numpy as np
12
  from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
13
  import requests
14
  from RealESRGAN import RealESRGAN
15
 
16
+
17
  import subprocess
18
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
19
 
 
115
 
116
  return enhanced_text
117
 
 
118
  def upscale_image(image, scale):
119
+ # Convert PIL Image to numpy array
120
+ img_np = np.array(image)
121
+
122
  if scale == 2:
123
+ upscaled_np = realesrgan_x2.predict(img_np)
124
  elif scale == 4:
125
+ upscaled_np = realesrgan_x4.predict(img_np)
126
  else:
127
  return image
128
+
129
+ # Convert numpy array back to PIL Image
130
+ return Image.fromarray(upscaled_np)
131
 
132
  @spaces.GPU(duration=120)
133
  def generate_image(additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps,