Geek7 commited on
Commit
3b57a34
·
verified ·
1 Parent(s): 528ac74

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -123
app.py CHANGED
@@ -1,123 +0,0 @@
1
- from flask import Flask, request, jsonify, send_file
2
- from flask_cors import CORS
3
- import os
4
- from huggingface_hub import InferenceClient
5
- from io import BytesIO
6
- from PIL import Image
7
-
8
- # Initialize the Flask app
9
- app = Flask(__name__)
10
- CORS(app) # Enable CORS for all routes
11
-
12
- # Initialize the InferenceClient with your Hugging Face token
13
- HF_TOKEN = os.environ.get("HF_TOKEN") # Ensure to set your Hugging Face token in the environment
14
- client = InferenceClient(token=HF_TOKEN)
15
-
16
- @app.route('/')
17
- def home():
18
- return "Welcome to the Image Background Remover!"
19
-
20
- # Simple content moderation function
21
- def is_prompt_explicit(prompt):
22
- explicit_keywords = ["sexual", "nudity", "erotic", "explicit", "porn", "pornographic", "xxx", "hentai", "fetish", "sex", "sensual", "nude", "strip", "stripping", "adult", "lewd", "provocative", "obscene", "vulgar", "intimacy", "intimate", "lust", "arouse", "seductive", "seduction", "kinky", "bdsm", "dominatrix", "bondage", "hardcore", "softcore", "topless", "bottomless", "threesome", "orgy", "incest", "taboo", "masturbation", "genital", "penis", "vagina", "breast", "boob", "nipple", "butt", "anal", "oral", "ejaculation", "climax", "moan", "foreplay", "intercourse", "naked", "exposed", "suicide", "self-harm", "overdose", "poison", "hang", "end life", "kill myself", "noose", "depression", "hopeless", "worthless", "die", "death", "harm myself"] # Add more keywords as needed
23
- for keyword in explicit_keywords:
24
- if keyword.lower() in prompt.lower():
25
- return True
26
- return False
27
-
28
- # Function to generate an image from a text prompt
29
- def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/sd-3.5", num_inference_steps=50, guidance_scale=7.5, seed=None):
30
- try:
31
- # Generate the image using Hugging Face's inference API with additional parameters
32
- image = client.text_to_image(
33
- prompt=prompt,
34
- negative_prompt=negative_prompt,
35
- height=height,
36
- width=width,
37
- model=model,
38
- num_inference_steps=num_inference_steps, # Control the number of inference steps
39
- guidance_scale=guidance_scale, # Control the guidance scale
40
- seed=seed # Control the seed for reproducibility
41
- )
42
- return image # Return the generated image
43
- except Exception as e:
44
- print(f"Error generating image: {str(e)}")
45
- return None
46
-
47
- # Function to refine an image using the refiner model
48
- def refine_image(image, prompt, negative_prompt=None, model="stabilityai/stable-diffusion-xl-refiner-1.0", num_inference_steps=15, guidance_scale=7.5):
49
- try:
50
- # Use Hugging Face's image-to-image API to refine the image
51
- refined_image = client.image_to_image(
52
- prompt=prompt,
53
- negative_prompt=negative_prompt,
54
- image=image,
55
- model=model,
56
- num_inference_steps=num_inference_steps,
57
- guidance_scale=guidance_scale
58
- )
59
- return refined_image
60
- except Exception as e:
61
- print(f"Error refining image: {str(e)}")
62
- return None
63
-
64
- @app.route('/generate_image', methods=['POST'])
65
- def generate_api():
66
- data = request.get_json()
67
-
68
- # Extract required fields from the request
69
- prompt = data.get('prompt', '')
70
- negative_prompt = data.get('negative_prompt', None)
71
- height = data.get('height', 1024) # Default height
72
- width = data.get('width', 720) # Default width
73
- num_inference_steps = data.get('num_inference_steps', 50) # Default number of inference steps
74
- guidance_scale = data.get('guidance_scale', 7.5) # Default guidance scale
75
- model_name = data.get('model', 'stabilityai/sd-3.5') # Base model
76
- refiner_model_name = 'stabilityai/sd-xl-refiner-1.0' # Refiner model
77
- seed = data.get('seed', None) # Seed for reproducibility, default is None
78
-
79
- if not prompt:
80
- return jsonify({"error": "Prompt is required"}), 400
81
-
82
- try:
83
- # Check for explicit content
84
- if is_prompt_explicit(prompt):
85
- # Return the pre-defined "thinkgood.png" image
86
- return send_file(
87
- "thinkgood.jpeg",
88
- mimetype='image/png',
89
- as_attachment=False,
90
- download_name='thinkgood.png'
91
- )
92
-
93
- # Step 1: Generate the base image
94
- base_image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
95
-
96
- if not base_image:
97
- return jsonify({"error": "Failed to generate base image"}), 500
98
-
99
- # Step 2: Refine the image with the refiner model
100
- refined_image = refine_image(base_image, prompt, negative_prompt, refiner_model_name, num_inference_steps, guidance_scale)
101
-
102
- if not refined_image:
103
- return jsonify({"error": "Failed to refine image"}), 500
104
-
105
- # Save the refined image to a BytesIO object
106
- img_byte_arr = BytesIO()
107
- refined_image.save(img_byte_arr, format='PNG') # Convert the image to PNG
108
- img_byte_arr.seek(0) # Move to the start of the byte stream
109
-
110
- # Send the refined image as a response
111
- return send_file(
112
- img_byte_arr,
113
- mimetype='image/png',
114
- as_attachment=False, # Send the file inline
115
- download_name='refined_image.png' # File name for download
116
- )
117
- except Exception as e:
118
- print(f"Error in generate_api: {str(e)}") # Log the error
119
- return jsonify({"error": str(e)}), 500
120
-
121
- # Add this block to make sure your app runs when called
122
- if __name__ == "__main__":
123
- app.run(host='0.0.0.0', port=7860) # Run directly if needed for testing