Update server.py
Browse files
server.py
CHANGED
@@ -7,8 +7,6 @@ import os
|
|
7 |
from PIL import Image
|
8 |
from diffusers import StableDiffusionPipeline
|
9 |
import os
|
10 |
-
from tinydb import TinyDB, Query
|
11 |
-
import uuid
|
12 |
|
13 |
token = os.getenv("HF_TOKEN")
|
14 |
|
@@ -35,9 +33,6 @@ CORS(app)
|
|
35 |
stable_diff_pipe = None
|
36 |
model = None
|
37 |
|
38 |
-
os.makedirs('/data', exist_ok=True)
|
39 |
-
db = TinyDB('/data/artwork_results.json')
|
40 |
-
|
41 |
def load_models(model_name="rupeshs/LCM-runwayml-stable-diffusion-v1-5"):
|
42 |
global stable_diff_pipe, model
|
43 |
|
@@ -65,13 +60,6 @@ def extract_image_features(image):
|
|
65 |
|
66 |
return generated_features
|
67 |
|
68 |
-
# Helper to get or set a unique user_id via cookies
|
69 |
-
def get_user_id():
|
70 |
-
user_id = request.cookies.get('user_id')
|
71 |
-
if not user_id:
|
72 |
-
user_id = str(uuid.uuid4())
|
73 |
-
return user_id
|
74 |
-
|
75 |
@app.route('/')
|
76 |
def index():
|
77 |
return render_template('index.html')
|
@@ -80,10 +68,6 @@ def index():
|
|
80 |
def check_membership():
|
81 |
try:
|
82 |
model_name = request.form.get('model', 'rupeshs/LCM-runwayml-stable-diffusion-v1-5')
|
83 |
-
contact_info = request.form.get('contact_info') # Optional, if frontend sends it
|
84 |
-
opt_in = request.form.get('opt_in', 'false').lower() == 'true'
|
85 |
-
user_id = get_user_id()
|
86 |
-
# Ensure models are loaded with the selected model
|
87 |
global stable_diff_pipe, model
|
88 |
if stable_diff_pipe is None or model is None:
|
89 |
load_models(model_name)
|
@@ -100,35 +84,12 @@ def check_membership():
|
|
100 |
output = model(processed_features)
|
101 |
probability = output.item()
|
102 |
predicted = int(output > 0.5)
|
103 |
-
|
104 |
-
if predicted == 1:
|
105 |
-
entry = {
|
106 |
-
'model': model_name,
|
107 |
-
'user_id': user_id,
|
108 |
-
'timestamp': int(uuid.uuid1().time),
|
109 |
-
}
|
110 |
-
if opt_in and contact_info:
|
111 |
-
entry['contact_info'] = contact_info
|
112 |
-
# Only store if not already present for this user/model
|
113 |
-
User = Query()
|
114 |
-
exists = db.search((User.model == model_name) & (User.user_id == user_id))
|
115 |
-
if not exists:
|
116 |
-
db.insert(entry)
|
117 |
-
# Count unique user_ids for this model
|
118 |
-
User = Query()
|
119 |
-
user_ids = set(r['user_id'] for r in db.search(User.model == model_name))
|
120 |
-
count = len(user_ids)
|
121 |
-
response = make_response(jsonify({
|
122 |
'probability': probability,
|
123 |
'predicted_class': predicted,
|
124 |
'message': f"Predicted membership probability: {probability}",
|
125 |
-
'is_in_training_data': "Likely" if predicted == 1 else "Unlikely"
|
126 |
-
|
127 |
-
}))
|
128 |
-
# Set user_id cookie if not present
|
129 |
-
if not request.cookies.get('user_id'):
|
130 |
-
response.set_cookie('user_id', user_id, max_age=60*60*24*365*5) # 5 years
|
131 |
-
return response
|
132 |
except Exception as e:
|
133 |
print(f"Error processing request: {str(e)}")
|
134 |
return jsonify({'error': str(e)}), 500
|
|
|
7 |
from PIL import Image
|
8 |
from diffusers import StableDiffusionPipeline
|
9 |
import os
|
|
|
|
|
10 |
|
11 |
token = os.getenv("HF_TOKEN")
|
12 |
|
|
|
33 |
stable_diff_pipe = None
|
34 |
model = None
|
35 |
|
|
|
|
|
|
|
36 |
def load_models(model_name="rupeshs/LCM-runwayml-stable-diffusion-v1-5"):
|
37 |
global stable_diff_pipe, model
|
38 |
|
|
|
60 |
|
61 |
return generated_features
|
62 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
63 |
@app.route('/')
|
64 |
def index():
|
65 |
return render_template('index.html')
|
|
|
68 |
def check_membership():
|
69 |
try:
|
70 |
model_name = request.form.get('model', 'rupeshs/LCM-runwayml-stable-diffusion-v1-5')
|
|
|
|
|
|
|
|
|
71 |
global stable_diff_pipe, model
|
72 |
if stable_diff_pipe is None or model is None:
|
73 |
load_models(model_name)
|
|
|
84 |
output = model(processed_features)
|
85 |
probability = output.item()
|
86 |
predicted = int(output > 0.5)
|
87 |
+
return jsonify({
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
88 |
'probability': probability,
|
89 |
'predicted_class': predicted,
|
90 |
'message': f"Predicted membership probability: {probability}",
|
91 |
+
'is_in_training_data': "Likely" if predicted == 1 else "Unlikely"
|
92 |
+
})
|
|
|
|
|
|
|
|
|
|
|
93 |
except Exception as e:
|
94 |
print(f"Error processing request: {str(e)}")
|
95 |
return jsonify({'error': str(e)}), 500
|