Update app.py

app.py

@@ -6,10 +6,15 @@ from indic_transliteration.detect import detect as detect_script
 from indic_transliteration.sanscript import transliterate
 import langdetect
 import re
+import requests
+import json
+import base64
+from PIL import Image
+import io
 
 # Initialize clients
 text_client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-image_client = InferenceClient("ijohn07/DALLE-4K")
+SPACE_URL = "https://ijohn07-dalle-4k.hf.space"
 
 def detect_language_script(text: str) -> tuple[str, str]:
     """Detect language and script of the input text.
@@ -111,20 +116,56 @@ def is_image_request(message: str) -> bool:
     message_lower = message.lower()
     return any(trigger in message_lower for trigger in image_triggers)
 
-def generate_image(prompt: str) -> str:
-    """Generate an image using DALLE-4K model."""
+def generate_image_space(prompt: str) -> Image.Image:
+    """Generate an image using the DALLE-4K Space."""
     try:
-        response = image_client.text_to_image(
-            prompt,
-            parameters={
-                "negative_prompt": "blurry, bad quality, nsfw",
-                "num_inference_steps": 30,
-                "guidance_scale": 7.5
-            }
-        )
-        # Save the image and return the path or base64 string
-        # Note: Implementation depends on how you want to handle the image output
-        return response
+        # First get the session hash
+        response = requests.post(f"{SPACE_URL}/queue/join")
+        session_hash = response.json().get('session_hash')
+        
+        # Send the generation request
+        payload = {
+            "prompt": prompt,
+            "negative_prompt": "blurry, bad quality, nsfw",
+            "num_inference_steps": 30,
+            "guidance_scale": 7.5,
+            "session_hash": session_hash
+        }
+        
+        response = requests.post(f"{SPACE_URL}/run/predict", json={
+            "data": [
+                prompt,  # Prompt
+                "",     # Negative prompt
+                7.5,    # Guidance scale
+                30,     # Steps
+                "DPM++ SDE Karras",  # Scheduler
+                False,  # High resolution
+                False,  # Image to image
+                None,   # Image upload
+                1       # Batch size
+            ],
+            "session_hash": session_hash
+        })
+        
+        # Poll for results
+        while True:
+            status_response = requests.post(f"{SPACE_URL}/queue/status", json={
+                "session_hash": session_hash
+            })
+            status_data = status_response.json()
+            
+            if status_data.get('status') == 'complete':
+                # Get the image data
+                image_data = status_data['data']['image']
+                # Convert base64 to PIL Image
+                image_bytes = base64.b64decode(image_data.split(',')[1])
+                image = Image.open(io.BytesIO(image_bytes))
+                return image
+            elif status_data.get('status') == 'error':
+                raise Exception(f"Image generation failed: {status_data.get('error')}")
+            
+            time.sleep(1)  # Wait before polling again
+            
     except Exception as e:
         print(f"Image generation error: {e}")
         return None
@@ -174,11 +215,9 @@ def respond(
     # Check if this is an image generation request
     if is_image_request(message):
         try:
-            image = generate_image(message)
+            image = generate_image_space(message)
             if image:
-                yield f"Here's your generated image based on: {message}"
-                # You'll need to implement the actual image display logic
-                # depending on your Gradio interface requirements
+                yield (image, f"Here's your generated image based on: {message}")
                 return
             else:
                 yield "Sorry, I couldn't generate the image. Please try again."
@@ -258,7 +297,4 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)"
         ),
     ]
-)
-
-if __name__ == "__main__":
-    demo.launch(share=True)
+)