Update app.py

app.py

@@ -1,71 +1,44 @@
 import gradio as gr
-import requests
-import os
-import time
-
-# Environment variables for API details
-API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")  # Fetching the API token from environment variable
-
-# Function to query Hugging Face API
-def query_huggingface_api(api_url, prompt):
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    data = {"inputs": prompt}
-    response = requests.post(api_url, headers=headers, json=data)
-    
-    if response.status_code == 200:
-        return response.content, None  # Return the image and no error
+from diffusers import StableDiffusionPipeline
+import torch
+
+# Function to automatically switch between GPU and CPU
+def load_model(model_id):
+    if torch.cuda.is_available():
+        device = "cuda"
+        info = "Running on GPU (CUDA)"
     else:
-        return None, f"Error {response.status_code}: {response.text}"  # Return None and the error message
-
-# Gradio function for generating the image
-def generate_image(api_url, prompt):
-    # Attempt to query the API with retry logic for loading models
-    for attempt in range(5):  # Try up to 5 times
-        result, error = query_huggingface_api(f"https://api-inference.huggingface.co/models/{api_url}", prompt)
-        
-        if result:
-            return result, None
-        elif "Model is currently loading" in error:
-            estimated_time = float(error.split("estimated_time\":")[1].split("}")[0])  # Extract estimated time from error message
-            time.sleep(estimated_time + 5)  # Wait for the model to load, with an additional buffer time
-        else:
-            return None, error  # Return the error if it's not a loading issue
+        device = "cpu"
+        info = "Running on CPU"
+    
+    # Load the model dynamically on the correct device
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+    pipe = pipe.to(device)
+    
+    return pipe, info
 
-    return None, "Model is still loading after multiple attempts. Please try again later."  # Final error if all attempts fail
+# Function for text-to-image generation with dynamic model ID and device info
+def generate_image(model_id, prompt):
+    pipe, info = load_model(model_id)
+    image = pipe(prompt).images[0]
+    return image, info
 
-# Create Gradio Blocks Interface
-with gr.Blocks(theme="nevreal/blues") as demo:
-    gr.Markdown(
-        """
-        # Text to Image Generator
-        Enter a text prompt, and the custom model will generate an image.
-        """
-    )
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Custom Text-to-Image Generator")
 
     with gr.Row():
         with gr.Column():
-            text_input = gr.Textbox(
-                label="Enter your prompt", 
-                placeholder="Type something here...",
-                value=""
-            )
-            model_input = gr.Textbox(
-                label="Model URL", 
-                placeholder="Enter the model URL...",
-                value=""
-            )
+            model_id = gr.Textbox(label="Enter Model ID (e.g., nevreal/vMurderDrones)", placeholder="Model ID")
+            prompt = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to generate")
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column():
-            image_output = gr.Image(label="Generated Image")
-            error_output = gr.Textbox(label="Error", interactive=False)
-
-    # Define the action for the button
-    generate_btn.click(
-        fn=generate_image, 
-        inputs=[model_input, text_input],  # Pass both model URL and prompt
-        outputs=[image_output, error_output]
-    )
+            output_image = gr.Image(label="Generated Image")
+            device_info = gr.Markdown()  # To display if GPU or CPU is used
+    
+    # Link the button to the image generation function
+    generate_btn.click(fn=generate_image, inputs=[model_id, prompt], outputs=[output_image, device_info])
 
-# Launch the Gradio Blocks WebUI
-demo.launch(share=True, debug=True)
+# Launch the app
+demo.launch()