nisten commited on
Commit
b8261fb
1 Parent(s): f3e62ee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -5
app.py CHANGED
@@ -1,9 +1,8 @@
1
- import os
2
- os.system('pip install transformers')
3
-
4
  import gradio as gr
5
  from transformers import OlmoeForCausalLM, AutoTokenizer
6
  import torch
 
 
7
 
8
  # Initialize ZeroGPU
9
  os.environ["ZEROGPU"] = "1"
@@ -28,6 +27,15 @@ def generate_text(prompt):
28
  out = model.generate(**inputs, max_length=64)
29
  return tokenizer.decode(out[0])
30
 
 
 
 
 
 
31
  # Set up the Gradio chat interface
32
- iface = gr.ChatInterface(fn=generate_text, system_prompt=system_prompt)
33
- iface.launch()
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
  from transformers import OlmoeForCausalLM, AutoTokenizer
3
  import torch
4
+ import os
5
+ from gradio_client import Client
6
 
7
  # Initialize ZeroGPU
8
  os.environ["ZEROGPU"] = "1"
 
27
  out = model.generate(**inputs, max_length=64)
28
  return tokenizer.decode(out[0])
29
 
30
+ # Function to set client for session
31
+ def set_client_for_session(request: gr.Request):
32
+ x_ip_token = request.headers['x-ip-token']
33
+ return Client("gradio/text-to-image", headers={"X-IP-Token": x_ip_token})
34
+
35
  # Set up the Gradio chat interface
36
+ with gr.Blocks() as demo:
37
+ client = gr.State()
38
+ iface = gr.ChatInterface(fn=generate_text, system_prompt=system_prompt)
39
+
40
+ demo.load(set_client_for_session, None, client)
41
+ iface.launch()