Anupam251272's picture
Update app.py
ea7fb29 verified
# Install required packages
import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
import os
from web3 import Web3
import json
from PIL import Image
import io
import ipfshttpclient
import requests
from pathlib import Path
# Function to check GPU availability
def get_device():
if torch.cuda.is_available():
return "cuda"
return "cpu"
# Initialize the model
def initialize_model():
device = get_device()
print(f"Using device: {device}")
model_id = "runwayml/stable-diffusion-v1-5"
pipeline = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
safety_checker=None
)
pipeline = pipeline.to(device)
return pipeline
# Initialize IPFS client
def initialize_ipfs():
try:
# Try to connect to local IPFS daemon
client = ipfshttpclient.connect('/ip4/127.0.0.1/tcp/5001')
return client
except Exception as e:
print(f"Could not connect to local IPFS: {e}")
# Fallback to Infura IPFS (you'll need to sign up for an Infura account)
return None
# Function to upload to IPFS (with fallback to Infura)
def upload_to_ipfs(image, ipfs_client=None):
try:
if ipfs_client:
# Convert PIL image to bytes
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
# Upload to IPFS
result = ipfs_client.add(img_byte_arr)
return result['Hash']
else:
# Fallback: Save locally
temp_path = "temp_artwork.png"
image.save(temp_path)
return f"Local file saved: {temp_path}"
except Exception as e:
return f"Error uploading to IPFS: {str(e)}"
# Initialize blockchain connection
def initialize_blockchain():
try:
# Connect to Polygon Mumbai testnet
w3 = Web3(Web3.HTTPProvider('https://rpc-mumbai.maticvigil.com'))
print(f"Connected to blockchain: {w3.is_connected()}")
return w3
except Exception as e:
print(f"Error connecting to blockchain: {e}")
return None
# Combined function to generate and process artwork
def generate_and_process_artwork(prompt, num_inference_steps, guidance_scale):
try:
# Generate the image
image = pipeline(
prompt,
num_inference_steps=int(num_inference_steps),
guidance_scale=float(guidance_scale)
).images[0]
# Try to upload to IPFS
ipfs_result = upload_to_ipfs(image, ipfs_client)
# Create metadata
metadata = {
"name": "AI Generated Artwork",
"description": f"Generated using Stable Diffusion with prompt: {prompt}",
"prompt": prompt,
"generation_params": {
"steps": num_inference_steps,
"guidance_scale": guidance_scale
}
}
success_message = f"Generation successful!\n"
if isinstance(ipfs_result, str) and ipfs_result.startswith("Qm"):
success_message += f"IPFS Hash: {ipfs_result}"
else:
success_message += "Artwork saved locally"
return image, success_message
except Exception as e:
return None, f"Error: {str(e)}"
# Create Gradio interface
def create_interface():
with gr.Blocks() as demo:
gr.Markdown("# AI Art Generation with Blockchain Integration")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(
label="Prompt",
placeholder="Enter your art description here...",
lines=3
)
steps = gr.Slider(
minimum=20,
maximum=100,
value=50,
step=1,
label="Number of Inference Steps"
)
guidance = gr.Slider(
minimum=1,
maximum=20,
value=7.5,
step=0.5,
label="Guidance Scale"
)
generate_btn = gr.Button("Generate Artwork")
with gr.Column():
output_image = gr.Image(label="Generated Artwork")
output_text = gr.Textbox(label="Status")
generate_btn.click(
fn=generate_and_process_artwork,
inputs=[prompt, steps, guidance],
outputs=[output_image, output_text]
)
gr.Markdown("""
## Features
- AI Art Generation using Stable Diffusion
- IPFS Storage Integration
- Automatic Metadata Generation
## Tips for Better Results
- Be specific in your prompts
- Use descriptive adjectives
- Include art style references
""")
return demo
# Initialize components
print("Initializing model...")
pipeline = initialize_model()
print("Initializing IPFS...")
ipfs_client = initialize_ipfs()
print("Initializing blockchain...")
w3 = initialize_blockchain()
# Launch the interface
demo = create_interface()
demo.launch(debug=True, share=True)