s-ahal commited on
Commit
0f03893
·
verified ·
1 Parent(s): 84eb1d6

Update server.py

Browse files
Files changed (1) hide show
  1. server.py +3 -42
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
- # Store in DB if likely and opt-in
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
- 'likely_count': count
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