Update app.py

app.py

@@ -1,51 +1,34 @@
 import gradio as gr
-import requests
+from transformers import StableDiffusionPipeline
+import torch
 from PIL import Image
-import base64
-from io import BytesIO
-
-def query_hf_image_generation(api_key, prompt):
-    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json"
-    }
-    data = {"inputs": prompt}
-
-    response = requests.post(API_URL, headers=headers, json=data)
-
-    if response.status_code != 200:
-        return f"Error: Received HTTP {response.status_code} - {response.text}"
-
-    try:
-        result = response.json()
-    except ValueError:
-        return f"Error decoding JSON: Unexpected response format {response.text}"
-
-    if 'error' in result:
-        return f"Error: {result['error']}"
-
-    if 'data' in result:
-        try:
-            base64_string = result['data'][0]
-            base64_data = base64_string.split(",")[1] if "," in base64_string else base64_string
-            image_data = base64.b64decode(base64_data)
-            image = Image.open(BytesIO(image_data))
-            return image
-        except Exception as e:
-            return f"Error processing image data: {e}"
-    else:
-        return "Error: Missing 'data' in the response."
-
-iface = gr.Interface(
-    fn=query_hf_image_generation,
-    inputs=[
-        gr.Textbox(label="Hugging Face API Key", placeholder="Enter your Hugging Face API Key here..."),
-        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
-    ],
-    outputs=gr.Image(label="Generated Image"),
-    title="Stable Diffusion XL Image Generator",
-    description="Enter your API Key and a prompt to generate an image using the Stable Diffusion XL model from Hugging Face."
-)
+import requests
 
-iface.launch(share=True)
+def generate_image(prompt):
+    # Load the preprocessing and model pipeline
+    # Here, we assume the Kvikontent/midjourney-v6 model has text-to-image capabilities in a manner similar to stable diffusion.
+    # This part needs verification and adjustment according to actual model documentation and availability.
+    model_id = "Kvikontent/midjourney-v6"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Setup the model pipeline (this can be adjusted if the model's actual interface differs)
+    # This example uses the typical usage pattern for generative models, but you should adjust according to the actual model's specs.
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)  # Replace with actual method to load Kvikontent/midjourney-v6 if different
+    pipe = pipe.to(device)
+    
+    # Generating the image
+    image = pipe(prompt).images[0]  # This line assumes the return type is accessible like this, adjust this according to actual usage.
+    
+    # Convert tensor to PIL Image (adjust if the output format differs)
+    image = Image.fromarray(image.numpy(), 'RGB')
+    return image
+
+# Create a Gradio interface
+iface = gr.Interface(fn=generate_image,
+                     inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
+                     outputs="image",
+                     title="Text to Image Generator",
+                     description="Type some text and generate an image using the Kvikontent/midjourney-v6 model.")
+
+# Running the application
+iface.launch()