Hugging Face's logo

Spaces:
Geek7
/
tadmztxi
Runtime error

App Files Community
Geek7 commited on
Commit
bbb7c5d
·
verified ·
1 Parent(s): 4713500

Update myapp.py

Browse files
Files changed (1) hide show
  1. myapp.py +27 -56
myapp.py CHANGED
@@ -1,71 +1,42 @@
1
  from flask import Flask, request, jsonify
2
- from diffusers import DiffusionPipeline
3
  import torch
4
- from PIL import Image
5
- import os
6
-
7
-
8
 
9
  # Initialize the Flask app
10
- myapp = Flask(__name__)
11
 
12
- # Load the Diffusion pipeline
13
- pipe = DiffusionPipeline.from_pretrained("prompthero/openjourney-v4").to("cpu")
 
14
 
15
- @myapp.route('/')
16
- def index():
17
- return '''
18
- <html>
19
- <body>
20
- <h1>Welcome to the Image Generation API!</h1>
21
- <form id="input-form">
22
- <label for="prompt">Enter your prompt:</label><br>
23
- <input type="text" id="prompt" name="prompt"><br><br>
24
- <button type="submit">Generate Image</button>
25
- </form>
26
- <div id="spinner" style="display:none;">Generating image, please wait...</div>
27
- <div id="result"></div>
28
- <script>
29
- document.getElementById('input-form').onsubmit = async (e) => {
30
- e.preventDefault();
31
- document.getElementById('spinner').style.display = 'block';
32
- const prompt = document.getElementById('prompt').value;
33
 
34
- const response = await fetch('/generate_image', {
35
- method: 'POST',
36
- headers: { 'Content-Type': 'application/json' },
37
- body: JSON.stringify({ prompt })
38
- });
39
 
40
- const data = await response.json();
41
- document.getElementById('spinner').style.display = 'none';
42
- if (response.ok) {
43
- document.getElementById('result').innerHTML = `<h2>Image Generated:</h2><img src="${data.image_path}" alt="Generated Image">`;
44
- } else {
45
- document.getElementById('result').innerText = 'Error generating image: ' + data.error;
46
- }
47
- };
48
- </script>
49
- </body>
50
- </html>
51
- '''
52
 
53
- @myapp.route('/generate_image', methods=['POST'])
54
- def generate_image():
55
  data = request.json
56
- prompt = data.get('prompt', 'Astronaut in a jungle, cold color palette, muted colors, detailed, 8k')
57
 
58
- # Generate the image
59
- image = pipe(prompt).images[0]
60
 
61
- # Convert to PIL Image and save
62
- pil_image = Image.fromarray(image.numpy())
63
- output_path = f"{prompt.replace(' ', '_')}.png" # Create a file name based on the prompt
64
- pil_image.save(output_path)
65
 
66
- # Return the path to the generated image
67
- return jsonify({'image_path': output_path})
68
 
69
  if __name__ == "__main__":
70
- # Set the host and port
71
- myapp.run(host='0.0.0.0', port=7860)
 
1
  from flask import Flask, request, jsonify
 
2
  import torch
3
+ from transformers import AutoModel, AutoTokenizer
4
+ from fastsafetensors import safe_load
 
 
5
 
6
  # Initialize the Flask app
7
+ app = Flask(__name__)
8
 
9
+ # Load the model and tokenizer using safe_load
10
+ model_path = "https://huggingface.co/prompthero/openjourney-v4/blob/main/safety_checker/model.safetensors" # Replace with your .safetensors file path
11
+ model_data = safe_load(model_path)
12
 
13
+ # Specify the model name, adjust as necessary
14
+ model_name = "prompthero/openjourney-v4" # Replace with your model name
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
 
16
+ # Load the tokenizer
17
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
 
18
 
19
+ # Load the model weights from safeload
20
+ model = AutoModel.from_pretrained(model_name, state_dict=model_data)
21
+
22
+ @app.route('/')
23
+ def index():
24
+ return "Welcome to the AI Model API!"
 
 
 
 
 
 
25
 
26
+ @app.route('/generate', methods=['POST'])
27
+ def generate_output():
28
  data = request.json
29
+ prompt = data.get('prompt', 'Hello, world!')
30
 
31
+ # Tokenize input prompt
32
+ inputs = tokenizer(prompt, return_tensors="pt")
33
 
34
+ # Generate output
35
+ with torch.no_grad():
36
+ outputs = model(**inputs)
 
37
 
38
+ # Process and return the output
39
+ return jsonify(outputs)
40
 
41
  if __name__ == "__main__":
42
+ app.run(host='0.0.0.0', port=5000)