From 9994bb67a1cc1a09f4c5eba70dfbd77c96d0ed76 Mon Sep 17 00:00:00 2001
From: hlohaus <983577+hlohaus@users.noreply.github.com>
Date: Wed, 5 Feb 2025 17:07:20 +0100
Subject: Add zerogpu_uuid to demo

---
 g4f/Provider/hf_space/G4F.py          | 18 +++++++++++++++---
 g4f/Provider/hf_space/Janus_Pro_7B.py |  5 ++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

(limited to 'g4f/Provider/hf_space')

diff --git a/g4f/Provider/hf_space/G4F.py b/g4f/Provider/hf_space/G4F.py
index 1cff0097..9f16d789 100644
--- a/g4f/Provider/hf_space/G4F.py
+++ b/g4f/Provider/hf_space/G4F.py
@@ -34,10 +34,21 @@ class G4F(Janus_Pro_7B):
         height: int = 1024,
         seed: int = None,
         cookies: dict = None,
+        zerogpu_token: str = None,
+        zerogpu_uuid: str = None,
         **kwargs
     ) -> AsyncResult:
         if cls.default_model not in model:
-            async for chunk in super().create_async_generator(model, messages, prompt=prompt, seed=seed, cookies=cookies, **kwargs):
+            async for chunk in super().create_async_generator(
+                model, messages,
+                proxy=proxy,
+                prompt=prompt,
+                seed=seed,
+                cookies=cookies, 
+                zerogpu_token=zerogpu_token,
+                zerogpu_uuid=zerogpu_uuid,
+                **kwargs
+            ):
                 yield chunk
             return
 
@@ -64,8 +75,9 @@ class G4F(Janus_Pro_7B):
             "trigger_id": 10
         }
         async with ClientSession() as session:
-            yield Reasoning(status="Acquiring GPU Token")
-            zerogpu_uuid, zerogpu_token = await get_zerogpu_token(cls.space, session, JsonConversation(), cookies)
+            if zerogpu_token is None:
+                yield Reasoning(status="Acquiring GPU Token")
+                zerogpu_uuid, zerogpu_token = await get_zerogpu_token(cls.space, session, JsonConversation(), cookies)
             headers = {
                 "x-zerogpu-token": zerogpu_token,
                 "x-zerogpu-uuid": zerogpu_uuid,
diff --git a/g4f/Provider/hf_space/Janus_Pro_7B.py b/g4f/Provider/hf_space/Janus_Pro_7B.py
index 0484316b..69708f13 100644
--- a/g4f/Provider/hf_space/Janus_Pro_7B.py
+++ b/g4f/Provider/hf_space/Janus_Pro_7B.py
@@ -70,6 +70,8 @@ class Janus_Pro_7B(AsyncGeneratorProvider, ProviderModelMixin):
         prompt: str = None,
         proxy: str = None,
         cookies: Cookies = None,
+        zerogpu_token: str = None,
+        zerogpu_uuid: str = None,
         return_conversation: bool = False,
         conversation: JsonConversation = None,
         seed: int = None,
@@ -90,7 +92,8 @@ class Janus_Pro_7B(AsyncGeneratorProvider, ProviderModelMixin):
         session_hash = generate_session_hash() if conversation is None else getattr(conversation, "session_hash")
         async with StreamSession(proxy=proxy, impersonate="chrome") as session:
             session_hash = generate_session_hash() if conversation is None else getattr(conversation, "session_hash")
-            zerogpu_uuid, zerogpu_token = await get_zerogpu_token(cls.space, session, conversation, cookies)
+            if zerogpu_token is None:
+                zerogpu_uuid, zerogpu_token = await get_zerogpu_token(cls.space, session, conversation, cookies)
             if conversation is None or not hasattr(conversation, "session_hash"):
                 conversation = JsonConversation(session_hash=session_hash, zerogpu_token=zerogpu_token, zerogpu_uuid=zerogpu_uuid)
             conversation.zerogpu_token = zerogpu_token
-- 
cgit v1.2.3