soiz commited on
Commit
d1c54ae
β€’
1 Parent(s): 9c05761

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -39
app.py CHANGED
@@ -1,10 +1,9 @@
1
  import os
2
- import gradio as gr
3
  from random import randint
4
- from diffusers import DiffusionPipeline
5
  from flask import Flask, request, send_file
6
  from io import BytesIO
7
  from PIL import Image, ImageChops
 
8
 
9
  app = Flask(__name__)
10
 
@@ -15,10 +14,7 @@ def load_model(model_name):
15
  global models_load
16
  if model_name not in models_load:
17
  try:
18
- # Assume models are hosted on the HuggingFace model hub
19
- model = DiffusionPipeline.from_pretrained(model_name)
20
- model = model.to("cuda")
21
- models_load[model_name] = model
22
  except Exception as error:
23
  print(f"Error loading model {model_name}: {error}")
24
  models_load[model_name] = None
@@ -27,46 +23,42 @@ def gen_fn(model_str, prompt, negative_prompt=None, noise=None, cfg_scale=None,
27
  if model_str not in models_load:
28
  load_model(model_str) # γƒ’γƒ‡γƒ«γŒγƒ­γƒΌγƒ‰γ•γ‚Œγ¦γ„γͺγ„ε ΄εˆγ―γƒ­γƒΌγƒ‰γ™γ‚‹
29
 
30
- if model_str in models_load:
31
- model = models_load[model_str]
32
- if model is None:
33
- return None, f"Model {model_str} could not be loaded."
34
 
35
- if noise == "random":
36
- noise = str(randint(0, 99999999999))
37
- full_prompt = f'{prompt} {noise}' if noise else prompt
38
 
39
- try:
40
- # Generate the image with the model
41
- if cfg_scale is not None and num_inference_steps is not None:
42
- result = model(full_prompt, negative_prompt=negative_prompt, cfg_scale=cfg_scale, num_inference_steps=num_inference_steps).images[0]
43
- elif cfg_scale is not None:
44
- result = model(full_prompt, negative_prompt=negative_prompt, cfg_scale=cfg_scale).images[0]
45
- elif num_inference_steps is not None:
46
- result = model(full_prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps).images[0]
47
- else:
48
- result = model(full_prompt, negative_prompt=negative_prompt).images[0]
 
49
 
50
- # Ensure the result is an image
51
- if isinstance(result, Image.Image):
52
- image = result
53
- else:
54
- print("Result is not an image:", type(result))
55
- return None, 'Result is not an image'
56
-
57
  # Check if the image is completely black
58
- black = Image.new('RGB', image.size, (0, 0, 0))
59
- if ImageChops.difference(image, black).getbbox() is None:
60
  return None, 'The image is completely black. There may be a parameter that cannot be specified, or an error may have occurred internally.'
 
61
 
62
- return image, None
 
 
63
 
64
- except Exception as e:
65
- print("Error generating image:", e)
66
- return None, f"Error generating image: {e}"
67
- else:
68
- print(f"Model {model_str} not found")
69
- return None, f"Model {model_str} not found"
70
 
71
  @app.route('/', methods=['GET'])
72
  def home():
 
1
  import os
 
2
  from random import randint
 
3
  from flask import Flask, request, send_file
4
  from io import BytesIO
5
  from PIL import Image, ImageChops
6
+ from diffusers import DiffusionPipeline
7
 
8
  app = Flask(__name__)
9
 
 
14
  global models_load
15
  if model_name not in models_load:
16
  try:
17
+ models_load[model_name] = DiffusionPipeline.from_pretrained(f'models/{model_name}')
 
 
 
18
  except Exception as error:
19
  print(f"Error loading model {model_name}: {error}")
20
  models_load[model_name] = None
 
23
  if model_str not in models_load:
24
  load_model(model_str) # γƒ’γƒ‡γƒ«γŒγƒ­γƒΌγƒ‰γ•γ‚Œγ¦γ„γͺγ„ε ΄εˆγ―γƒ­γƒΌγƒ‰γ™γ‚‹
25
 
26
+ model = models_load.get(model_str)
27
+ if not model:
28
+ print(f"Model {model_str} not found or failed to load.")
29
+ return None, f"Model {model_str} not found or failed to load."
30
 
31
+ if noise == "random":
32
+ noise = str(randint(0, 99999999999))
33
+ full_prompt = f'{prompt} {noise}' if noise else prompt
34
 
35
+ try:
36
+ # Generate the image
37
+ kwargs = {"prompt": full_prompt}
38
+ if negative_prompt:
39
+ kwargs["negative_prompt"] = negative_prompt
40
+ if cfg_scale is not None:
41
+ kwargs["guidance_scale"] = cfg_scale
42
+ if num_inference_steps is not None:
43
+ kwargs["num_inference_steps"] = num_inference_steps
44
+
45
+ result = model(**kwargs).images[0]
46
 
47
+ # Check if result is an image
48
+ if isinstance(result, Image.Image):
 
 
 
 
 
49
  # Check if the image is completely black
50
+ black = Image.new('RGB', result.size, (0, 0, 0))
51
+ if ImageChops.difference(result, black).getbbox() is None:
52
  return None, 'The image is completely black. There may be a parameter that cannot be specified, or an error may have occurred internally.'
53
+ return result, None
54
 
55
+ else:
56
+ print("Result is not an image:", type(result))
57
+ return None, 'Result is not an image'
58
 
59
+ except Exception as e:
60
+ print("Error generating image:", e)
61
+ return None, f"Error generating image: {e}"
 
 
 
62
 
63
  @app.route('/', methods=['GET'])
64
  def home():